卷积神经网络LeNet-5的pytorch代码实现

在详解卷积神经网络LeNet-5一文中，我详细介绍了一下Lenet-5的基本原理，下面简要介绍一下Lenet-5的pytorch代码实现。
建议本文对应上文一起查看。

一、Lenet-5网络模型实现

LeNet-5 神经网络一共五层，其中卷积层和池化层可以考虑为一个整体，网络的结构为 ：

输入 → 卷积 → 池化 → 卷积 → 池化 → 卷积（全连接） → 全连接 → 输出。
编程之前，先看一下相关数值：

层数	in_channel	out_channel	kernel_size	stride	padding
卷积层C1	1	6	5	1	2
池化层S2	6	6	2	2	0
卷积层C3	6	16	5	1	0
池化层S4	16	16	2	2	0

【注】：池化层的in_channel和out_channel不会变。

层数	输入	输出
卷积层C5	16	120
全连接层F6	120	84
输出层	84	10

Step1： 引入必要的包

代码如下：

import torch
from torch import nn
import torch.nn.functional as F

Step2：搭建卷积层C1和池化层S2

代码如下：

		#输入:(6*28*28);输出：(6*14*14)
        self.conv1 = nn.Sequential(     
            nn.Conv2d(1,6,kernel_size=5,stride=1,padding=2), #新建卷积层,kernel_size表示卷积核大小,stride表示步长,padding=2，图片大小变为 28+2*2 = 32 (两边各加2列0)，保证输入输出尺寸相同
            nn.ReLU(),
            #可以选择取最大值池化MaxPool2d,也可以选择取平均值池化AvgPool2d,两者参数相同
            nn.MaxPool2d(kernel_size = 2 ,stride = 2,padding=0)   #input_size=(6*28*28)，output_size=(6*14*14)
        )

这段代码对应于第1个卷积层和第1个池化层。

nn.Conv2d是2D卷积的意思，in_channels定义输入通道的数量，out_channels定义输出通道的数量，kernel_size定义卷积核的尺寸，stride定义卷积步长，padding表示增加的边框尺寸。
输出的size尺寸：

size = (h - kernel_size + 2*padding)/stride + 1

接着每个激活函数后面连接一个最大池化（现在用的比较多的也是最大池化），这样，我们这两层神经网络就搭建完成了。

Step3：搭建卷积层C3和池化层S4

代码如下：

#输入:(6*14*14);输出：(16*5*5)
self.conv2 = nn.Sequential(
    nn.Conv2d(6,16,kernel_size=5,stride=1,padding=0), #input_size=(6*14*14)，output_size=16*10*10
    nn.ReLU(),
    nn.MaxPool2d(kernel_size = 2,stride = 2,padding=0)    ##input_size=(16*10*10)，output_size=(16*5*5)
)

这段代码对应于第2个卷积层和第2个池化层。

Step4：搭建全连接层C5、全连接层F6以及输出层

代码如下：

#全连接层
self.fc1 = nn.Sequential(
    nn.Linear(16*5*5,120),
    nn.ReLU()
)
#全连接层
self.fc2 = nn.Sequential(
    nn.Linear(120,84),
    nn.ReLU()
)
#全连接层
self.fc3 = nn.Linear(84,10)

Step5：设置网络前向传播

#网络前向传播过程
def forward(self,x):
    x = self.conv1(x)
    x = self.conv2(x)
 
    x = x.view(x.size(0), -1) #全连接层均使用的nn.Linear()线性结构，输入输出维度均为一维，故需要把数据拉为一维
    x = self.fc1(x)
    x = self.fc2(x)
    x = self.fc3(x)
    return x

Step6：查看神经网络

代码如下：

#查看网络
myNet = LeNet5()
print(myNet) 

输出结果：

二、Lenet-5实现MNIST手写数字识别

Step1：载入训练集和测试集

代码：

#load data
transform = torchvision.transforms.ToTensor()   #定义数据预处理方式：转换 PIL.Image 成 torch.FloatTensor
 
train_data = torchvision.datasets.MNIST(root="E:\\python_mnist\mnist\\train",    #数据目录，这里目录结构要注意。
                                        train=True,                                     #是否为训练集
                                        transform=transform,                            #加载数据预处理
                                        download=False)                                 #是否下载
test_data = torchvision.datasets.MNIST(root="E:\\python_mnist\mnist\\train",
                                        train=False,
                                        transform=transform,
                                        download=False)
#数据加载器:组合数据集和采样器 
train_loader = torch.utils.data.DataLoader(dataset = train_data,batch_size = 64,shuffle = True) 
test_loader = torch.utils.data.DataLoader(dataset = test_data,batch_size = 64,shuffle = False)

第一次训练时，就把参数download改成True，先下载再训练。

Step2：观察部分MNIST数据集

代码如下：

import numpy as np
import matplotlib.pyplot as plt
def imshow(img):
     img = img / 2 + 0.5 # unnormalize
     npimg = img.numpy()
     plt.imshow(np.transpose(npimg, (1, 2, 0)))
     plt.show()

# torchvision.utils.make_grid 将图片进行拼接
imshow(torchvision.utils.make_grid(iter(train_loader).next()[0]))

输出结果：

Step3：设置损失函数和优化器

代码如下：

#define loss
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")   #若检测到GPU环境则使用GPU，否则使用CPU
net = LeNet().to(device)    #实例化网络，有GPU则将网络放入GPU加速
loss_fuc = nn.CrossEntropyLoss()    #多分类问题，选择交叉熵损失函数
optimizer = optim.SGD(net.parameters(),lr = 0.001,momentum = 0.9)   #选择SGD，学习率取0.001

Step4：开始训练

代码如下：

#开始训练
EPOCH = 8   #迭代次数
for epoch in range(EPOCH):
    sum_loss = 0
    #数据读取
    for i,data in enumerate(train_loader):
        inputs,labels = data
        inputs, labels = inputs.to(device), labels.to(device)   #有GPU则将数据置入GPU加速
 
        # 梯度清零
        optimizer.zero_grad()
 
        # 传递损失 + 更新参数
        output = net(inputs)
        loss = loss_fuc(output,labels)
        loss.backward()
        optimizer.step()
 
        # 每训练100个batch打印一次平均loss
        sum_loss += loss.item()
        if i % 100 == 99:
            print('[Epoch:%d, batch:%d] train loss: %.03f' % (epoch + 1, i + 1, sum_loss / 100))
            sum_loss = 0.0
 
    correct = 0
    total = 0
 
    for data in test_loader:
        test_inputs, labels = data
        test_inputs, labels = test_inputs.to(device), labels.to(device)
        outputs_test = net(test_inputs)
        _, predicted = torch.max(outputs_test.data, 1)  #输出得分最高的类
        total += labels.size(0) #统计50个batch 图片的总个数
        correct += (predicted == labels).sum()  #统计50个batch 正确分类的个数
 
    print('第{}个epoch的识别准确率为：{}%'.format (epoch + 1, 100*correct.item()/total))

#模型保存
torch.save(net.state_dict(),'D:\\Desktop\\ckpt.mdl') 

#模型加载
#net.load_state_dict(torch.load('ckpt.mdl'))

运行结果：

Step5：当然，也可以设置loss曲线

加入代码：

loss_list = []
......
# 打印损失值变化曲线
import matplotlib.pyplot as plt
plt.plot(loss_list)
plt.title('traning loss')
plt.xlabel('epochs')
plt.ylabel('loss')
plt.show()

结果：

关于详细代码，链接在这里：
https://download.csdn.net/download/didi_ya/12805831

需要的自取。

如果觉得有帮助，记得点个赞哟~