pytorch实战1：手把手教你基于pytorch实现LeNet5

前言

最近在看经典的卷积网络架构，打算自己尝试复现一下，在此系列文章中，会参考很多文章，有些已经忘记了出处，所以就不贴链接了，希望大家理解。

完整的代码在最后。

本系列必须的基础

python基础知识、CNN原理知识、pytorch基础知识

本系列的目的

一是帮助自己巩固知识点；

二是自己实现一次，可以发现很多之前的不足；

三是希望可以给大家一个参考。

目录结构

文章目录

- pytorch实战1：手把手教你基于pytorch实现LeNet5
- - 1. 数据集介绍与下载：
  - - 1.1 介绍：
    - 1.2 下载：
  - 2. LeNet5模型创建：
  - - 2.1 架构介绍：
    - 2.2 模型类创建：
  - 3. 模型训练和评估：
  - - 3.1 数据加载：
    - 3.2 模型实例化和放入GPU中：
    - 3.3 定义损失函数、优化器：
    - 3.4 训练：
    - 3.5 评估：
    - 3.6 探究：
  - 4. 总结：

1. 数据集介绍与下载：

1.1 介绍：

LeNet5是1998年提出的，主要用来当时的手写数字识别，因此使用的数据集是MNIST数据集。

MNIST是一个经典的手写数字数据，也是一个公开的小型数据。 MNIST中的图像每个都是28*28=784的大小，并且为灰度图，值为0-255。

1.2 下载：

数据集可以通过官网进行下载

http://yann.lecun.com/exdb/mnist/

，（建议）也可以通过

pytorch代码获取

：

# 导包
from torchvision.datasets import MNIST
import torchvision.transforms as transforms
# 下载数据集或者加载数据集
train_dataset = MNIST(root='../data/',train=True,transform=transforms.ToTensor(),download=True)
test_dataset = MNIST(root='../data/',train=False,transform=transforms.ToTensor())

注意，上面下载代码中root参数指明了下载后保存的地址，需要根据自己的文件夹路径进行修改。

2. LeNet5模型创建：

2.1 架构介绍：

LeNet5模型架构如下图所示：

pytorch实战1：手把手教你基于pytorch实现LeNet5

其中，值得注意的地方是：

输入大小为32*32，而数据集的图像大小为28*28，因此需要做出处理，方法是为原图像进行填充，即padding=2（这样大小为：28+2+2=32）。
上图中全连接层写的不是特别清晰，比如FC1层120个神经元是指的是输出的神经元个数，输入神经元个数应该是5*5*16。

2.2 模型类创建：

下面我们来创建这个模型类，首先根据pytorch创建模型的基本结构，写出：

# 创建模型
class LeNet(nn.Module):
    def __init__(self):
        super(LeNet,self).__init__()
        # 定义模型
        pass

    def forward(self,x):
		pass

接着，我们来定义我们的模型，这里我们采取的思路是：先定义前面的卷积、池化层，再定义全连接层。

好的，先定义卷积层（看注释）：

self.features = nn.Sequential(
    # C1层，输入通道数为1是因为为灰度图，不是彩色图。其余的就是根据架构填写的参数，除去padding=2，
    # padding=2是为了让28*28的图片变为32*32
    nn.Conv2d(in_channels=1,out_channels=6,kernel_size=(5,5),stride=1,padding=2),
    nn.ReLU(),
    # Pool 1 层
    nn.MaxPool2d(kernel_size=2,stride=2),
    # C2层，输入通道数是上一层的输出通道数，其余同上
    nn.Conv2d(in_channels=6,out_channels=16,kernel_size=(5,5),stride=1),
    nn.ReLU(),
    # Pool 2 层
    nn.MaxPool2d(kernel_size=2,stride=2),
)

这里，我把本来使用的sigmoid函数改为了relu函数，大家可以写的时候可以改回来。

然后，定义全连接层（这里注意，我自己写的时候开始还有点懵，后来才理解）：

self.classifier = nn.Sequential(
    # FC1层，输入为5*5*16=400，输出数为指定的参数值
    nn.Linear(in_features=400, out_features=120),
    nn.ReLU(),
    # FC2 层
    nn.Linear(in_features=120, out_features=84),
    nn.ReLU(),
    # FC3 层
    nn.Linear(in_features=84, out_features=10)
)

最开始写的时候，没有想到还有400这个值，所以很纠结，后来才想起来：卷积层到全连接层的时候，需要把多维的数据拉平，所以显然此时的输入为卷积层输出的各个维度相乘，即6*6*25。

最后，来定义前向算法，这个非常简单，只是注意要多一步把数据拉平（变为1维）的操作：

def forward(self,x):
    # 定义前向算法
    x = self.features(x)
    x = torch.flatten(x,1)
    result = self.classifier(x)
    return result

这里，也许你会问：**为什么flatten参数要写一个1？这是因为，我们这里会采取批量训练，因此传入的数据是一个类似于[batch,32,32,5]的思维数据，其中batch指的是每批的个数，后面分别书图像大小（32-32）和卷积核个数。因此，我们拉平的时候，需要从第二位开始拉平，使之变为[batch,32*32*5]**的形式，因此需要指定为1。

3. 模型训练和评估：

3.1 数据加载：

使用pytorch加载数据，很简单：

# 下载数据集或者加载数据集
train_dataset = MNIST(root='../data/',train=True,transform=transforms.ToTensor(),download=True)
test_dataset = MNIST(root='../data/',train=False,transform=transforms.ToTensor())
# 加载数据: 分批次，每批32个数据
batch_size = 32
train_loader = DataLoader(dataset=train_dataset,batch_size=batch_size,shuffle=True)
test_loader = DataLoader(dataset=test_dataset,batch_size=batch_size,shuffle=False)

**基本上（也有需要自己定义Dataset的情况）**加载已有的数据都是这样的格式。

3.2 模型实例化和放入GPU中：

这段很简单，另外调用GPU的方法也是固定的，就不多说：

# 创建模型
model = LeNet()
# 模型放入GPU中
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)

3.3 定义损失函数、优化器：

损失函数我们使用分类任务中常用的交叉熵损失函数，而优化器采用SGD优化器，学习率设置维常用的0.001，Momentum设置为0.9。

# 定义损失函数
loss_func = nn.CrossEntropyLoss()
loss_list = [] # 用来存储损失值
# 定义优化器：第一个参数是模型的参数
SGD = optim.SGD(params=model.parameters(),lr=0.001,momentum=0.9)

3.4 训练：

首先，指定训练次数：

# 训练指定次数，这里写为了3
for i in range(3):
	xx

然后，迭代读取数据：

for i in range(3):
    loss_temp = 0 # 定义一个损失值，用来打印查看
    # 其中j是迭代次数，data和label都是批量的，每批32个
    for j,(batch_data,batch_label) in enumerate(train_loader):
    	xxx

接着将数据放入GPU中：

for i in range(3):
    loss_temp = 0
    for j,(batch_data,batch_label) in enumerate(train_loader):
        # 启用GPU
        batch_data,batch_label = batch_data.cuda(),batch_label.cuda()

然后，就是常见的操作：清空梯度、计算模型、计算损失、反向传播、更新梯度

for i in range(3):
    loss_temp = 0 # 定义一个损失值，用来打印查看
    # 其中j是迭代次数，data和label都是批量的，每批32个
    for j,(batch_data,batch_label) in enumerate(train_loader):
        # 启用GPU
        batch_data,batch_label = batch_data.cuda(),batch_label.cuda()
        # 清空梯度
        SGD.zero_grad()
        # 模型训练
        prediction = model(batch_data)
        # 计算损失
        loss = loss_func(prediction,batch_label)
        loss_temp += loss
        # BP算法
        loss.backward()
        # 更新梯度
        SGD.step()

最后，我们可以打印一下损失值来查看模型训练的怎么样了，这里我们每隔两百小批次就打印一次损失值：

# 训练指定次数
for i in range(3):
    loss_temp = 0 # 定义一个损失值，用来打印查看
    # 其中j是迭代次数，data和label都是批量的，每批32个
    for j,(batch_data,batch_label) in enumerate(train_loader):
        # 启用GPU
        batch_data,batch_label = batch_data.cuda(),batch_label.cuda()
        # 清空梯度
        SGD.zero_grad()
        # 模型训练
        prediction = model(batch_data)
        # 计算损失
        loss = loss_func(prediction,batch_label)
        loss_temp += loss
        # BP算法
        loss.backward()
        # 更新梯度
        SGD.step()
        if (j + 1) % 200 == 0:
            print('第%d次训练，第%d批次，损失值: %.3f' % (i + 1, j + 1, loss_temp / 200))
            loss_temp = 0

运行上面的代码，显示的结果为：

第1次训练，第200批次，损失值: 2.298
第1次训练，第400批次，损失值: 2.285
第1次训练，第600批次，损失值: 2.241
第1次训练，第800批次，损失值: 1.812
第1次训练，第1000批次，损失值: 0.754
第1次训练，第1200批次，损失值: 0.521
第1次训练，第1400批次，损失值: 0.423
第1次训练，第1600批次，损失值: 0.365
第1次训练，第1800批次，损失值: 0.334
第2次训练，第200批次，损失值: 0.284
第2次训练，第400批次，损失值: 0.245
第2次训练，第600批次，损失值: 0.235
第2次训练，第800批次，损失值: 0.211
第2次训练，第1000批次，损失值: 0.202
第2次训练，第1200批次，损失值: 0.179
第2次训练，第1400批次，损失值: 0.180
第2次训练，第1600批次，损失值: 0.165
第2次训练，第1800批次，损失值: 0.148
第3次训练，第200批次，损失值: 0.147
第3次训练，第400批次，损失值: 0.143
第3次训练，第600批次，损失值: 0.136
第3次训练，第800批次，损失值: 0.130
第3次训练，第1000批次，损失值: 0.115
第3次训练，第1200批次，损失值: 0.114
第3次训练，第1400批次，损失值: 0.114
第3次训练，第1600批次，损失值: 0.093
第3次训练，第1800批次，损失值: 0.122

3.5 评估：

我们按照上面训练的思路，可以轻松写出测试的代码：

correct = 0
for batch_data,batch_label in test_loader:
    batch_data, batch_label = batch_data.cuda(), batch_label.cuda()
    prediction = model(batch_data)
    predicted = torch.max(prediction.data, 1)[1]
    correct += (predicted == batch_label).sum()
print('准确率: %.2f %%' % (100 * correct / 10000)) # 因为总共10000个测试数据

当然上面的代码第一次接触还是有点疑问，主要的一点是

predicted = torch.max(prediction.data, 1)[1]

这段代码在干什么。

首先，我们一个图像传入模型，输出的是一个向量，这个向量10个值，代表数字0-9的概念值。而torch.max(x,1)表示按行（1：行，0：列）取出最大值，它返回的是一个特殊格式的数据，第一个元素是各个最大值，第二个元素是其索引（在这里就等价于0-9数字），因此使用[1]取出。

运行结果：

准确率: 91.25 %
# 这个是只训练一次的结果

3.6 探究：

下面将模型训练1、2、3、4次的准确率结果：

# 1次 ： 准确率: 91.25 %
# 2次 ： 准确率: 93.44 %
# 3次 ： 准确率: 97.19 %
# 4次 ： 准确率: 97.58 %

另外，测试一下使用GPU与不使用GPU的时间差：

# 为了测试，我训练10次，并且仅仅记录训练花费时间
# 使用GPU： 训练花了: 124 s
# 不适用GPU：训练花了: 160 s

上面GPU测试不严谨，因为首次调用GPU是需要花费时间的，但是从两者差别看出，调用GPU确实效率很好。

4. 总结：

这次算是把整个流程从头到尾跑了一遍，并简单探究了训练次数和调用GPU对训练的影响。会了本篇文章的代码，至少后面再实现CNN的图像分类架构，应该还是比较简单了。

完整代码

# author: baiCai
# 导包
import time
import torch
from torch import nn
from torch import optim
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
import torchvision.transforms as transforms
from torchvision.models import AlexNet

# 创建模型
class LeNet(nn.Module):
    def __init__(self):
        super(LeNet,self).__init__()
        # 定义模型
        self.features = nn.Sequential(
            nn.Conv2d(in_channels=1,out_channels=6,kernel_size=(5,5),stride=1,padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2,stride=2),
            nn.Conv2d(in_channels=6,out_channels=16,kernel_size=(5,5),stride=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2,stride=2),
        )
        self.classifier = nn.Sequential(
            nn.Linear(in_features=400, out_features=120),
            nn.ReLU(),
            nn.Linear(in_features=120, out_features=84),
            nn.ReLU(),
            nn.Linear(in_features=84, out_features=10)
        )

    def forward(self,x):
        # 定义前向算法
        x = self.features(x)
        # print(x.shape)
        x = torch.flatten(x,1)
        # print(x.shape)
        result = self.classifier(x)
        return result

# 下载数据集或者加载数据集
train_dataset = MNIST(root='../data/',train=True,transform=transforms.ToTensor(),download=True)
test_dataset = MNIST(root='../data/',train=False,transform=transforms.ToTensor())
# 加载数据: 分批次，每批256个数据
batch_size = 32
train_loader = DataLoader(dataset=train_dataset,batch_size=batch_size,shuffle=True)
test_loader = DataLoader(dataset=test_dataset,batch_size=batch_size,shuffle=False)
# start time
start_time = time.time()
# 创建模型
model = LeNet()
# 模型放入GPU中
# device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# model = model.to(device)
# 定义损失函数
loss_func = nn.CrossEntropyLoss()
loss_list = [] # 用来存储损失值
# 定义优化器
SGD = optim.SGD(params=model.parameters(),lr=0.001,momentum=0.9)
# 训练指定次数
for i in range(10):
    loss_temp = 0 # 定义一个损失值，用来打印查看
    # 其中j是迭代次数，data和label都是批量的，每批32个
    for j,(batch_data,batch_label) in enumerate(train_loader):
        # 启用GPU
        # batch_data,batch_label = batch_data.cuda(),batch_label.cuda()
        # 清空梯度
        SGD.zero_grad()
        # 模型训练
        prediction = model(batch_data)
        # 计算损失
        loss = loss_func(prediction,batch_label)
        loss_temp += loss
        # BP算法
        loss.backward()
        # 更新梯度
        SGD.step()
        if (j + 1) % 200 == 0:
            print('第%d次训练，第%d批次，损失值: %.3f' % (i + 1, j + 1, loss_temp / 200))
            loss_temp = 0
# end_time
end_time = time.time()
print('训练花了: %d s' % int((end_time-start_time)))
# 使用GPU： 训练花了: 124 s
# 不适用GPU：训练花了: 160 s
# 测试
correct = 0
for batch_data,batch_label in test_loader:
    batch_data, batch_label = batch_data.cuda(), batch_label.cuda()
    prediction = model(batch_data)
    predicted = torch.max(prediction.data, 1)[1]
    correct += (predicted == batch_label).sum()
print('准确率: %.2f %%' % (100 * correct / 10000)) # 因为总共10000个测试数据

pytorch实战1：手把手教你基于pytorch实现LeNet5