【他山之石】从零开始实现一个卷积神经网络

“他山之石，可以攻玉”，站在巨人的肩膀才能看得更高，走得更远。在科研的道路上，更需借助东风才能更快前行。为此，我们特别搜集整理了一些实用的代码链接，数据集，软件，编程技巧等，开辟“他山之石”专栏，助你乘风破浪，一路奋勇向前，敬请关注。

作者：知乎—073-759

地址：https://zhuanlan.zhihu.com/p/355527103

本教程基于pytorch，使用LeNet作为神经网络模型，MNIST作为训练和测试的数据集。本教程将从python、pytorch、CUDA的安装开始，最后实现一个可用的手写数字分类器。本教程适合对python、深度学习和卷积神经网络有初步的了解。

01

准备工作

python环境搭建

首先需要进行python环境的搭建，可以去python官网[1]下载，也可以使用Anaconda[2]。一般推荐使用Anaconda，安装后即可直接使用，而且其自带的包管理器和非常方便。

CUDA下载及安装

如果你的电脑中有Nvidia的独立显卡，那么你可以使用CUDA技术对神经网络进行加速；若你的电脑中没有Nvidia的独立显卡，可直接跳过这一步。

在设备管理器-显示适配器中查看自己显卡的型号，并前往Nvidia显卡算力表[3]查询自己显卡的算力，当前最新版本的pytorch支持算力大于3.5的显卡，若你的显卡算力小于3.5，则需要安装老版本的pytorch和CUDA或者使用CPU进行模型的训练。

前往CUDA Toolkit[4]下载并安装CUDA，在此期间，你需要记住你所安装的CUDA版本号。安装完毕后，在命令行输入nvcc -V查看已安装CUDA版本号，若出现如下显示，则说明CUDA安装成功。

nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2019 NVIDIA Corporation

pytorch安装

进入pytorch官网[5]，点击install按钮进入下载界面[6]。PyTorch Build栏选择稳定版；Your OS栏选择你所用的操作系统；Package建议选择pip；Language选择python；Compute Platform栏选择你所安装的CUDA版本，若你的电脑不支持CUDA则选择None。最后，Run this Command栏会生成一行指令，示例如下：

pip install torch==1.8.0+cu111 torchvision==0.9.0+cu111 torchaudio===0.8.0 -f https://download.pytorch.org/whl/torch_stable.html

将这行命令复制粘贴至命令行后运行即可进行下载。若下载速度缓慢则可以考虑挂梯子或者更换国内镜像源，更换源方法详见链接[7]。

编辑器安装

python是一种解释型语言，理论上说任何编辑器或IDE都可以使用。一般来说，写python比较常用的编辑器是Visual Studio Code[8]，它可以写任何语言，是一个轻量化的编辑器；常用的IDE是Pycharm[9]，这是一个专用于python开发的IDE，拥有许多Visual Studio Code所不具备的功能，若你要进行项目的开发，Pycharm会更适合你。

02

建立工作文件夹

一般来说，将一个项目的代码、数据和文档放在一个文件夹内是一个良好的习惯，这方便使用相对路径进行调用，且移动的时候可以不用考虑路径的变化。因此，我们可以新建一个文件夹，并改名为CNN_MNIST，这样我们可以把接下来用到的代码都放在这个文件夹内，以方便管理。

03

构建模型

分析模型

建立好工作文件夹后，我们就可以开始搭建我们所用的模型了。我们需要的模型是LeNet-5，如下图所示。

从图中可以看出，其输入32x32的灰度图像，由于MNIST数据集的图像为28x28，因此，我们将输入改为28x28，并依次计算每一层输出的特征图大小。其每一层参数大致如下：

输入层：输入大小28x28，通道数为1。注意：本层不算LeNet-5的网络结构，一般情况下不将输入层视为网络层次结构之一。

C1-卷积层：输入大小28x28，通道数为1；输出大小28x28，通道数为6；卷积核大小为5x5；步长为1；边缘补零为2；激活函数为ReLU。注意：为了提升卷积神经网络的效果，在每个卷积层后添加激活函数，本教程使用的激活函数为ReLU。

S2-池化层：输入大小28x28，通道数为6；输出大小14x14，通道数为6；池化核大小为2x2；步长为2；池化方式为最大池化。

C3-卷积层：输入大小14x14，通道数为6；输出大小10x10，通道数为16；卷积核大小为5x5；步长为1；边缘补零为0；激活函数为ReLU。

S4-池化层：输入大小10x10，通道数为16；输出大小5x5，通道数为16；池化核大小为2x2；步长为2；池化方式为最大池化。

C5-卷积层：输入大小5x5，通道数为16；输出大小1x1，通道数为120；卷积核大小为5x5；步长为1；边缘补零为0；激活函数为ReLU。注意：这层也可以看作全连接层，可以通过全连接的方法实现。

F6-全连接层：输入为120维向量；输出为84维向量；激活函数为ReLU。

OUTPUT-输出层：输入为84维向量；输出为10维向量。注意：该层也是全连接层，且不带激活函数。

对于网络的具体分析及卷积、池化、全连接原理详见：LeNet-5详解[10]

搭建模型

分析完模型后，我们就可以开始搭建我们模型了。在工作文件夹内新建model.py文件，然后用编辑器打开文件开始定义模型。

首先，我们需要导入相关的包，主要是torch包。其中，torch.nn包含了自定义网络的基类，卷积层、池化层、激活函数、全连接层等构建，损失函数的调用等函数和方法，因此我们可以单独调用并重命名以方便使用。该部分写法如下：

import torch
import torch.nn as nn

导入完相关的包后我们就需要定义我们的模型了。在这里，我们把模型命名为LeNet，因此我们创建一个名为LeNet的类，该类继承了nn.Module类，写法如下：

class LeNet(nn.Module):

模型被定义后就需要对模型进行初始化和前向传播，在nn.Module中需要__init__和forward两个函数对一个模型进行实现。__init__函数即初始化，主要用于定义每一层的构成，如卷积、池化层等；forward函数即前向传播，主要用于确定每一层之间的顺序，使得模型可以正常使用。这就跟我们制作手链一样，在初始化部分中确定我们的手链中要用哪些珠子，这些珠子分别是怎么样的；在前向传播过程中用线将每一个珠子穿起来，使得其成为一个完整的手链。

初始化部分写法如下(注意，python注重缩进，请按照文中的代码正确缩进)：

    def __init__(self):
        super(LeNet, self).__init__()

之后我们就要定义我们的每一层了，在这个模型中我们有2维卷积层(nn.Conv2d)、ReLU激活函数(nn.ReLU)、最大池化层(nn.MaxPool2d)和全连接层(nn.Linear)，其写法和主要参数分别如下：

2维卷积层：

nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias, padding_mode)

in_channels：输入的通道数

out_channels：输出的通道数

kernel_size：卷积核的大小。若卷积核是方形，则只需要一个整数边长；若不是方形，则需要输入一个元组表示高和宽

stride：卷积核每次滑动的步长，默认为1

padding：设置边缘补零的大小（也就是在输入特征图外围增加几圈0）

dilation：控制卷积核之间的间距，默认为0；若使用空洞卷积则需要对该参数进行设置

groups：控制输入和输出之间的连接，平时不常用，若使用分组卷积则需要设置该参数

bias：是否设置偏置，默认为 True

padding_mode：边缘补零模式，默认为”zeros“

一般情况下，我们使用in_channels, out_channels, kernel_size, stride, padding这几个参数

官方文档：Pytorch-Conv2d[11]

最大池化层：

nn.MaxPool2d(kernel_size, stride, padding, dilation, return_indices, ceil_mode)

kernel_size：池化核的大小。

stride：池化核每次滑动的步长，默认与池化核大小一致

padding：设置边缘补零的大小（也就是在输入特征图外围增加几圈0）

dilation：控制窗口中元素步幅的参数

return_indices：表示返回值中是否包含最大值位置的索引，默认为False

ceil_mode：其用于计算输出特征图形状的时候，是使用向上取整还是向下取整。False表示向下取整，True表示向上取整，默认为False

一般情况下，我们使用kernel_size, stride这几个参数

官方文档：Pytorch-MaxPool2d[12]

ReLU激活函数：

nn.ReLU(inplace)

inplace：是否进行就地运算，就地运算可以节省内存，但是会使得输入数据发生改变，默认为False

一般情况下，我们不需要对ReLU设置参数

官方文档：Pytorch-ReLU[13]

全连接层：

nn.Linear(in_features, out_features, bias)

in_features：输入向量维度

out_features：输出向量维度

bias：是否设置偏置，默认为True

一般情况下，我们使用in_features, out_features这几个参数

官方文档：Pytorch-Linear[14]

知道以上层的设置方法后，我们就可以定义我们模型中的每一层了。写法如下：

        self.C1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5, stride=1, padding=2)
        self.R1 = nn.ReLU()
        self.S2 = nn.MaxPool2d(kernel_size=2)
        self.C3 = nn.Conv2d(6, 16, 5, 1, 0)
        self.R2 = nn.ReLU()
        self.S4 = nn.MaxPool2d(2)
        self.C5 = nn.Conv2d(16, 120, 5, 1, 0)
        self.R3 = nn.ReLU()
        self.F6 = nn.Linear(in_features=120, out_features=84)
        self.R4 = nn.ReLU()
        self.OUT = nn.Linear(84, 10)

以上就是每一层的定义方法。我们在完成定义之后就需要完成前向传播部分了。前向传播函数需要传入self和输入的变量，一般写为x，即forward(self, x)。然后我们在函数内把之前定义好的层按顺序调用，每一层在计算后会返回结果，我们需要一个变量进行保存，即c1 = self.C1(x)，在最后，我们需要将最后一步的计算结果返回。当网络中不存在跳跃连接或密集连接等分支结构的情况下，我们可以直接用x作为中间变量。

在pytorch中，我们的图像数据以一个四维的张量传入模型，其形状为[batch, channels, h, w]。其中，batch即批大小，我们一般会一次性将一批图像送进网络处理，这一批图像的数量即为批大小；channel即通道数，也就是之前卷积层的channels；h和w分别代表图像的高和宽。Conv2d和MaxPool2d都接受以上形状的输入，ReLU接受任意形状的输入，而Linear只接受传入一个二维的张量，形状为[batch, length]，length表示长度，即向量的维度。在这里，我们需要把之前卷积层输出的四维张量转换为二维张量，而.view()可以实现这个操作，我们在需要处理的张量上直接使用view方法，然后输入需要改变的维度，比如说我们最后一个卷积层生成的特征图形状为[batch, 120, 1, 1]，我们要将其转换为[batch, 120]，若已知batch的大小，我们就可以直接batch和120填入括号中。但是，在实际使用中batch大小可能会随着超参数的变化而改变，因此我们可以直接使用.size()方法，在括号内填上维度即可返回所在维度的大小，如x是一个形状为[16, 3, 384, 256]的张量，则x.size(0)为16，即为batch的大小。之后我们需要在view中填入另一个参数，在这里我们知道是120，就可以直接填写120，不过我们在这边也可以填入-1，填写-1会让电脑自动帮我们计算这一栏所需参数的大小，这个方法在实际搭建模型的时候非常好用，因此大家一般都会写为-1。最后就需要一个变量保存这部分的返回值了，我们依然可以直接用x来保存。综上，我们只需要x = x.view(x.size(0), -1)这一句话即可。

前向传播函数写法如下：

   def forward(self, x):
        x = self.C1(x)
        x = self.R1(x)
        x = self.S2(x)
        x = self.C3(x)
        x = self.R2(x)
        x = self.S4(x)
        x = self.C5(x)
        x = self.R3(x)
        x = x.view(x.size(0), -1)
        x = self.F6(x)
        x = self.R4(x)
        x = self.OUT(x)
        return x

至此，我们就把网络模型搭建好了。为了测试我们的模型是否可用，我们给这个文件设置一个入口，即if __name__ == "__main__":。这个部分内的代码会在运行当前文件时执行，而若把该文件作为一个包导入的时候则不会被执行，我们可以通俗地理解为C/C++的主函数。然后，我们要定义一下需要使用的模型，即model = LeNet()。定义完之后我们可以用print(model)把模型的结构打印出来以便于检查模型是否正确。之后需要定义一个张量传给模型，看一下模型能否正常工作。我们定义一个随机张量a = torch.randn(1, 1, 28, 28)，其形状为[1, 1, 28, 28]。把a传入模型中，并将其返回值赋给b，即b = model(a)。这时候我们就获得了模型运行后的结果，我们可以把b打印出来，或者打印b.size()查看其维度是否正确。若一切正确，那么恭喜你，你的第一个模型搭建成功了！

model.py文件完整代码如下：

import torch
import torch.nn as nn

class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.C1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5, stride=1, padding=2)
        self.R1 = nn.ReLU()
        self.S2 = nn.MaxPool2d(kernel_size=2)
        self.C3 = nn.Conv2d(6, 16, 5, 1, 0)
        self.R2 = nn.ReLU()
        self.S4 = nn.MaxPool2d(2)
        self.C5 = nn.Conv2d(16, 120, 5, 1, 0)
        self.R3 = nn.ReLU()
        self.F6 = nn.Linear(in_features=120, out_features=84)
        self.R4 = nn.ReLU()
        self.OUT = nn.Linear(84, 10)
    def forward(self, x):
        x = self.C1(x)
        x = self.R1(x)
        x = self.S2(x)
        x = self.C3(x)
        x = self.R2(x)
        x = self.S4(x)
        x = self.C5(x)
        x = self.R3(x)
        x = x.view(x.size(0), -1)
        x = self.F6(x)
        x = self.R4(x)
        x = self.OUT(x)
        return x

if __name__ == "__main__":
    model = LeNet()
    a = torch.randn(1, 1, 28, 28)
    b = model(a)
    print(b)

04

下载数据集

在大多数教程中，数据集的下载工作都是放在训练文件中的，但是数据集仅需要一次下载即可，若放在训练文件夹中则需要判断数据集是否存在，然后对相应的参数进行修改，相对来说比较麻烦。为此，我们可以在工作文件夹内新建一个download_dataset.py文件专门用于下载数据集。

在此之前，我们需要在工作文件夹内新建一个文件夹，为了方便起见，我们将其命名为data。我们的数据集将下载保存在这个文件夹中。

打开download_dataset.py文件，然后导入torchvision包，在torchvision.dataset中有很多经典的数据集，我们可以将其下载下来。我们需要下载MNIST数据集，在这里可以使用一句话来完成：torchvision.datasets.MNIST(root, train, transform, target_transform, download)，我们只需要设置root为我们的data文件夹，将download设置为True即可完成下载。

download_dataset.py文件完整代码如下：

import torchvision
torchvision.datasets.MNIST('./data', download=True)

如果下载速度缓慢或下载失败，则需要挂梯子或更换网络环境。若仍然无法下载，可以使用百度云手写数据集下载[15]。从百度云下载后将rar文件解压到data文件夹内即可使用。

05

训练模型

初始化和导入模型

现在，我们要开始训练我们的模型了。在工作文件夹内新建train.py文件，为了方便起见，我们依次导入如下包并重命名：

import torch
import torchvision
import torch.nn as nn
import torch.utils.data as Data

之后，我们需要从model.py文件中把我们搭建的模型文件导入并定义，即：

from model import LeNet
model = LeNet()

定义超参数、数据集和DataLoader

导入完成后，我们需要对一些超参数进行定义，如遍历轮次(Epoch)、批大小(batch_size)、初始学习率(lr)等。在这里，我们对数据集遍历5轮，批大小设置为64，初始学习率为0.001，即：

Epoch = 5
batch_size = 64
lr = 0.001

然后，我们可以导入我们下载的数据集了。MNIST包含训练集和测试集，在这里，我们只需要其训练集。因此，我们可以定义一个train_data用于导入MNIST的训练集，并利用torchvision.transforms.ToTensor()将形状为[h, w, channel]，值为0~255之间的uint8图像转换成形状为[channel, h ,w]，值在0~1之间的torch.FloatTensor：

train_data = torchvision.datasets.MNIST(root='./data/', train=True, transform=torchvision.transforms.ToTensor(), download=False)

定义完训练集后我们需要定义一个DataLoader将train_data中的数据喂给模型。在pytorch中，DataLoader的定义如下：

Data.DataLoader(dataset, batch_size, shuffle, sampler, batch_sampler, num_workers, collate_fn, pin_memory, drop_last, timeout, worker_init_fn, prefetch_factor, persistent_workers)

以下是其常用参数的介绍：

dataset：数据集，可以使用Data.DataSet类或者torchvision.datasets

batch_size：批大小，每次迭代送入模型的图像数量

shuffle：是否打乱数据集，默认为False

num_workers：使用的线程数，DataLoader支持多线程读取数据以提升效率，该值为0或1是使用单线程进行读取。一般情况下该值不要超过cpu的最大线程，如果使用GPU训练模型的话该值越大其显存占用也会越大，日常使用中需要根据电脑的配置进行调节。默认为0

pin_memory：是否使用锁页内存，可以理解为是否将数据集全部强制加载进内存，且不与虚拟内存进行交换，设为True的话可以使得模型的训练快一些，默认为False

drop_last：是否丢弃最后不足batch_size的数据。有时候，数据集并不能整除batch_size，最后一批图像的数量会小于batch_size，这个参数决定是否将这一批数据丢弃，默认为False

官方文档：Pytorch-Data[16]

了解DataLoader定义后我们就可以定义一个train_loader用于将数据集分批次传送给模型。其中，我们把之前定义好的train_data和batch_size传入；然后设置shuffle为True以打乱数据集；num_workers设为0，虽然这样做效率较低，但是拥有较好的兼容性，有时候多线程会出一些莫名其妙的错误，在后面可以适当增加这个的数值；将drop_last设置为True，丢弃最后一个批次，确保每个批次的批大小严格相同。train_loader写法如下：

train_loader = Data.DataLoader(train_data, batch_size=batch_size, shuffle=True, num_workers=0, drop_last=True)

定义损失函数和优化器

在完成以上步骤后，我们就可以定义我们需要的损失函数和优化器了。手写数字分类器是一个单标签分类问题，即每次预测结果的标签仅有一个，因此，我们可以使用交叉熵损失作为我们的损失函数，即nn.CrossEntropyLoss()。优化器我们使用当前常用的Adam优化器，其定义为torch.optim.Adam(parameters, lr)，我们需要传入模型的参数和初始学习率，模型参数我们可以使用model.parameters()来获得，初始学习率即我们先前定义好的lr。该部分写法如下：

loss_function = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=lr)

启用梯度

我们知道，神经网络在训练时需要计算梯度，并在反向传播时候使用，因此我们要启用pytorch的自动求导机制，并利用其实现反向传播和参数更新。在这里，我们只需要简单的一句torch.set_grad_enabled(True)即可实现，其表示在接下来的计算中每一次运算产生的节点都是可以求导的。

然后，我们需要使用model.train()方法，该方法用于启用Batch Normalization层和Dropout层。虽然在我们的模型中并没有这两层，但是我们不妨将其加上，并作为一个习惯，以免在我们真正需要时忘记。

该部分写法如下：

torch.set_grad_enabled(True)
model.train()

使用CUDA加速

神经网络的训练一般是在GPU上进行的，我们将模型转移至显卡只需要一句model.cuda()即可，训练时我们也需要将我们的输入数据x、标签y分别使用.cuda()传至显卡。这是大部分教程中常用的方法。但是，我们可以通过torch.device()来指定使用的设备device，然后通过.to()方法将模型和数据放到指定的设备上，这样我们就可以通过定义device来指定是在cpu还是显卡上进行训练了，而且在多显卡的情况下也可以指定使用其中的某一张显卡进行训练。

但是，并不是所有的设备都支持CUDA加速的，而torch.cuda.is_available()可以判断本设备是否支持CUDA，如果支持就返回True，不支持就返回False。有了这个函数，我们就可以让其自动判断是否支持CUDA加速并自动选择设备了。而不支持CUDA加速的设备我们可以使用cpu来进行。

该部分写法如下：

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)

训练

接下来就到了最重要的部分了——训练。

在pytorch中，神经网络的训练一般是分以下几个步骤进行的：

1) 获得DataLoader中的数据x和标签y

2) 将优化器的梯度清零

3) 将数据送入模型中获得预测的结果y_pred

4) 将标签和预测结果送入损失函数获得损失

5) 将损失值反向传播

6) 使用优化器对模型的参数进行更新

以上这六个步骤分别对应着代码中的六行，在pytorch中，只需要这六行即可完成一次迭代。但是，我们的数据集不仅仅有这一次迭代，而且我们也要遍历数据集不止一次。因此，我们首先需要一个循环用来遍历数据集，即Epoch：

for epoch in range(Epoch):

在写好第一个循环后我们确定了遍历数据集的次数，然后，我们就要写每一次遍历中的迭代了。DataLoader本质上是一个迭代器，因此我们在写这个循环的时候需要使用枚举enumerate()的方法，在这个循环中，我们对每一次迭代进行实现，即step：

    for step, data in enumerate(train_loader):

在写好这两个循环后，我们就要写上面的那六个步骤了：

1) 将data中的数据和标签取出，其中数据为x，标签为y。即x, y = data

2) 通过调用optimizer.zero_grad()将优化器的梯度清零

3) 将数据送入模型获得预测结果y_pred，在这里我们需要对数据进行处理，将其传至之前定义好的设备上，即y_pred = model(x.to(device, torch.float))

4) 将标签和预测结果送入损失函数获得损失。由于损失函数的计算是在之前定义好的设备上的，因此我们需要将y也传至设备，同时，y_pred的类型为torch.float，但是其下标的类型为torch.long，损失函数需要用到其下标，因此我们需要把y的类型变化为与y_pred的下标一致，即loss = loss_function(y_pred, y.to(device, torch.long))

5) 将损失值反向传播，这边直接调用loss.backward()即可

6) 使用优化器对模型的参数进行更新，在这边也只需要调用optimizer.step()即可

以上就是整个训练部分，pytorch的语法和python相近，因此，如果有python基础的话上手pytorch是十分容易的。

该部分完整写法如下：

for epoch in range(Epoch):
    for step, data in enumerate(train_loader):
        x, y = data
        optimizer.zero_grad()
        y_pred = model(x.to(device, torch.float))
        loss = loss_function(y_pred, y.to(device, torch.long))
        loss.backward()
        optimizer.step()

保存模型

训练结束后，我们需要把模型保存，否则在之后的测试和使用中只能重新训练并调用，这是非常麻烦的。因此，将模型保存下来是一个非常必要的步骤。在pytorch中有两种方法用来保存模型，即保存整个模型和保存模型参数，这两方法都是通过torch.save()实现的。

在本教程中，我们保存整个模型，这样可以把模型的定义和参数全保存在一个文件中，在接下来的测试中不需要导入并调用模型的定义文件，在不修改模型结构的前提下这种方法是较为方便的。我们把模型保存在工作文件夹的根目录中，并命名为LeNet.pkl。把我们训练好的模型model和路径放入torch.save()中即可。写法如下：

torch.save(model, './LeNet.pkl')

训练过程可视化

至此，我们成功地让我们的模型跑起来了。但是，当运行train.py文件后它只会生成一个模型文件，而训练过程中的损失和准确率我们却不得而知。因此，我们需要使得训练的过程可以被我们看见，以确定模型是否正确地被训练了。

我们可以在每隔一定的step后输出当前损失和准确率的平均值。MNIST的训练集共有六万张图像，而我们的batch_size是64且丢弃最后一批，因此在每个Epoch中有937个step，实际训练59968张图像。我们可以每迭代100次后输出当前Epoch的损失和准确率的平均值，并输出当前处在哪一次Epoch和step。

首先，我们先在训练部分第一个循环内定义损失函数的平均值和平均准确率，这边的0.0表示这两个变量类型为浮点型：

for epoch in range(Epoch):
    running_loss = 0.0
    acc = 0.0
    for step, data in enumerate(train_loader):

然后，我们需要对每次计算产生的损失进行相加，把结果放在running_loss中，因此，我们需要在反向传播后添加一个累加操作。由于loss在我们之前定义的设备上，因此我们需要获得loss的值，然后将其传回cpu并转换为float类型，即：

running_loss += float(loss.data.cpu())

在累加了损失后我们需要计算其准确率，在这里，我们的准确率可以使用预测正确的数量除以总数，因此，我们需要判断其预测的标签是多少。y_pred是一个二维的张量，其形状为[batch, channel]，在这边channel是10，即十个数字。如果我们将batch中的任意一行提取出来就获得了一个10维的向量，向量里的每个数代表与其下标所对应的标签的相关性，相关性越大则代表越有可能是这个数字。因此，我们需要获得这个向量中最大数的下标，在pytorch中，我们可以用.argmax(dim)方法实现，输入维度dim，即可返回这个维度下最大值的下标，即pred = y_pred.argmax(dim=1)。在此基础上，我们就可以计算其预测正确的数量了，先获取pred的值，然后传回cpu，用==判断是否相等，然后相加即可：

acc += (pred.data.cpu() == y.data).sum()

之后，在训练部分的末尾添加判断，判断当前step是否是该轮中的第100个，由于python的计数是从0开始的，所以第一百个step实际上是99。这里只需要用if和求模运算即可实现：

if step % 100 == 99:

实现判断后我们就需要将平均值给输出出来了，首先计算损失的平均值，即running_loss / step，由于这边的step是从0开始的，所以需要加1，即：

loss_avg = running_loss / (step + 1)

然后计算准确率，我们刚才计算了预测正确的数量，将其除以当前已预测图像的总数即可，总数量可以通过step * batch_size求得，然后将其转换为float类型，即：

acc_avg = float(acc / ((step + 1) * batch_size))

计算完以上数据后我们就可以依次将其输出了，使用print来输出这些数据：

print('Epoch', epoch + 1, ',step', step + 1, '| Loss_avg: %.4f' % loss_avg, '|Acc_avg:%.4f' % acc_avg)

如果你照着写到了这里，那么你已经可以很好地对模型进行训练了，而且可以在训练过程中直接查看损失和准确率的变化，当你训练完毕并获得LeNet.pkl文件后，恭喜你，你的第一个模型训练成功了！

train.py文件完整代码如下：

import torch
import torchvision
import torch.nn as nn
import torch.utils.data as Data
from model import LeNet

model = LeNet()
Epoch = 5
batch_size = 64
lr = 0.001
train_data = torchvision.datasets.MNIST(root='./data/', train=True, transform=torchvision.transforms.ToTensor(), download=False)
train_loader = Data.DataLoader(train_data, batch_size=batch_size, shuffle=True, num_workers=0, drop_last=True)
loss_function = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
torch.set_grad_enabled(True)
model.train()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)

for epoch in range(Epoch):
    running_loss = 0.0
    acc = 0.0
    for step, data in enumerate(train_loader):
        x, y = data
        optimizer.zero_grad()
        y_pred = model(x.to(device, torch.float))
        loss = loss_function(y_pred, y.to(device, torch.long))
        loss.backward()
        running_loss += float(loss.data.cpu())
        pred = y_pred.argmax(dim=1)
        acc += (pred.data.cpu() == y.data).sum()
        optimizer.step()
        if step % 100 == 99:
            loss_avg = running_loss / (step + 1)
            acc_avg = float(acc / ((step + 1) * batch_size))
            print('Epoch', epoch + 1, ',step', step + 1, '| Loss_avg: %.4f' % loss_avg, '|Acc_avg:%.4f' % acc_avg)

torch.save(model, './LeNet.pkl')

06

测试模型

初始化、导入模型和数据集

现在，我们完成了模型的训练并已经获得了一个可用的模型，这个时候，我们就可以对我们的模型进行测试了。在工作文件夹内新建test.py，依次导入以下包：

import torch
import torchvision
import torch.utils.data as Data

然后，我们就可以定义我们的测试集和DataLoader了。这里，我们选择MNIST的测试集用于模型的测试部分，导入方法和之前导入训练集是一样的，唯一的不同就是要把train设为False。而在DataLoader部分，我们不需要打乱数据集，然后一批次只需要送入一张图像，因此我们需要对DataLoader进行一些修改。写法如下：

test_data = torchvision.datasets.MNIST(root='./data/', train=False, transform=torchvision.transforms.ToTensor(), download=False)
test_loader = Data.DataLoader(test_data, batch_size=1, shuffle=False)

之后，定义我们需要使用的设备：

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

接下来，我们需要载入我们已经训练好的模型。在这里我们需要使用torch.load(f, map_location, pickle_module, pickle_load_args)函数。f为我们模型文件的路径和名称，在这里是'./LeNet.pkl'；map_location会重新映射使用的设备，一般情况下这个参数不需要任何的修改，但是如果你想要把一个用GPU训练的模型放在一个只有cpu的设备上时会发生一些错误，而这时就需要定义该参数了，我们可以在这里填上torch.device(device)以避免这个错误的发生。其他的参数基本上不需要进行任何的设置。写法如下：

net = torch.load('./LeNet.pkl',map_location=torch.device(device))

然后将我们的模型传至相应的设备即可：net.to(device)

至此，我们设置好了需要的测试集、DataLoader和之前训练的模型，并将模型传到了接下来需要使用的设备上。

关闭梯度

在测试阶段，我们不需要对模型的参数进行更新，因此我们可以关闭自动求导功能，并使用net.eval()方法屏蔽Dropout层、冻结BN层的参数，防止在测试阶段BN层发生参数更新，即：

torch.set_grad_enabled(False)
net.eval()

测试及输出结果

接下来，我们就要正式地测试我们的模型了。

首先，我们要获取测试集的大小，用于最后准确率的计算。我们可以先获取其数据，然后计算数据的大小来实现。这里，我们采用.size()方法，获取其第0维度的大小：

length = test_data.data.size(0)

接着，我们要定义准确率，用于统计我们模型在测试集上的准确率：acc = 0.0。

然后，我们需要遍历一次数据集，因此我们只需要一个for循环即可：

for i, data in enumerate(test_loader):

测试的部分与训练类似，也是将数据输入模型，然后获得输出，只不过不需要计算损失、反向传播、参数更新等步骤了。这里，我们只需要两步即可：

    x, y = data
    y_pred = net(x.to(device, torch.float))

现在，我们获得了预测的结果，我们需要获得预测到的标签，并计算预测正确的数量。这里与训练部分是一致的：

    pred = y_pred.argmax(dim=1)
    acc += (pred.data.cpu() == y.data).sum()

然后，我们在每一次预测后输出其预测的结果和对应的真实值即可，在输出时，我们需要将这两者转换为整型。即：

 print('Predict:', int(pred.data.cpu()), '|Ground Truth:', int(y.data))

最后，我们需要计算模型在测试集上的准确率，即用预测正确的个数除以测试集的大小，并将其输出出来。为了美观，我们可以把它写成百分比的形式，写法如下：

acc = (acc / length) * 100
print('Accuracy: %.2f' %acc, '%')

至此，模型测试完成。如果你照着写到了这里，并获得了你所训练地模型的准确度，那么，你现在已经学会了pytorch的基本操作，并能够独立的训练和测试一个模型了。恭喜你，你成功入门了pytorch！

test.py文件完整代码如下：

import torch
import torchvision
import torch.utils.data as Data

test_data = torchvision.datasets.MNIST(root='./data/', train=False, transform=torchvision.transforms.ToTensor(), download=False)
test_loader = Data.DataLoader(test_data, batch_size=1, shuffle=False)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net = torch.load('./LeNet.pkl',map_location=torch.device(device))
net.to(device)
torch.set_grad_enabled(False)
net.eval()
length = test_data.data.size(0)
acc = 0.0

for i, data in enumerate(test_loader):
    x, y = data
    y_pred = net(x.to(device, torch.float))
    pred = y_pred.argmax(dim=1)
    acc += (pred.data.cpu() == y.data).sum()
    print('Predict:', int(pred.data.cpu()), '|Ground Truth:', int(y.data))
acc = (acc / length) * 100
print('Accuracy: %.2f' %acc, '%')

07

结语

本教程查阅、参考了大量的文档和资料，耗时近百小时写成，制作不易，请多点赞、收藏。

源代码已上传至GitHub：https://github.com/ZnsNgk/CNN_Course

链接：

[1] https://www.python.org/

[2] https://www.anaconda.com/

[3] https://developer.nvidia.com/cuda-gpus#collapseOne

[4] https://developer.nvidia.com/cuda-toolkit

[5] https://pytorch.org/

[6] https://pytorch.org/get-started/locally/

[7] https://jingyan.baidu.com/article/d5c4b52b21b63e9b570dc574.html

[8] https://code.visualstudio.com/

[9] https://www.jetbrains.com/pycharm/

[10] https://cuijiahua.com/blog/2018/01/dl_3.html

[11] https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html#torch.nn.Conv2d

[12] https://pytorch.org/docs/stable/generated/torch.nn.MaxPool2d.html#torch.nn.MaxPool2d

[13] https://pytorch.org/docs/stable/generated/torch.nn.ReLU.html#torch.nn.ReLU

[14] https://pytorch.org/docs/stable/generated/torch.nn.Linear.html#torch.nn.Linear

[15] https://link.zhihu.com/?target=https%3A//pan.baidu.com/s/1YWfieeG1c8w4JkpXBfl-rQ，分享码hp2q

[16] https://pytorch.org/docs/stable/data.html#loading-batched-and-non-batched-data

本文目的在于学术交流，并不代表本公众号赞同其观点或对其内容真实性负责，版权归原作者所有，如有侵权请告知删除。

“他山之石”历史文章