在之前的文章中,我们使用pytorch下载并转换了MNIST的数据集,今天我们实现LeNet-5网络,使用下载的训练集对网络进行训练,然后使用测试集测试最终的效果。
MNIST数据集详解及可视化处理(pytorch)
首先,我们先看一下LeNet-5,LeNet-5是由“深度学习三巨头”之一、图灵奖得主Yann LeCun在一篇名为"Gradient-Based Learning Applied to Document Recognition"的paper(paper下载地址:https://www.researchgate.net/publication/2985446_Gradient-\Based_Learning_Applied_to_Document_Recognition )中提出的神经网络结构,在手写数字和机器打印字符上十分高效。
LeNet-5共有7层,不包含输入,每层都包含可训练参数;每个层有多个Feature Map,每个FeatureMap通过一种卷积滤波器提取输入的一种特征。
它是一个简单的前馈神经网络,它接受一个输入,然后一层接着一层地传递,最后输出计算的结果。
我们可以对LeNet-5的每一层分别做一个分析:
1.首先输入图像是单通道的3232大小的图像,用矩阵表示就是[1,32,32]。
2. c1卷积层,conv1所用的卷积核尺寸为55,滑动步长为1,卷积核数目为6,那么经过该层后图像尺寸变为28,32-5+1=28,输出矩阵为[6,28,28]。
3. s2池化层,pool核尺寸为22,步长2,这是没有重叠的max pooling,池化操作后,图像尺寸减半,变为14×14,输出矩阵为[6,14,14]。
4. c3卷积层,conv2的卷积核尺寸为55,步长1,卷积核数目为16,卷积后图像尺寸变为10,这是因为14-5+1=10,输出矩阵为[16,10,10]。
5. s4池化层,pool核尺寸为22,步长2,这是没有重叠的max pooling,池化操作后,图像尺寸减半,变为5×5,输出矩阵为[16,5,5]。
6. c5卷积层,使用55的卷积核对s4的输出进行卷积,卷积核数目为120,输出图片大小为5-5+1=1,最终输出120个11的特征图。这里实际上是与S4全连接了,但仍将其标为卷积层,原因是如果LeNet-5的输入图片尺寸变大,其他保持不变,那该层特征图的维数也会大于11。
7. F6全连接层,与C5层全连接,神经元个数为84。
8. output输出层,与F6层全连接,输出长度为10的张量,代表所输入的图像属于哪个类别。
现在我们在pytorch框架上实现LeNet-5网络,由于MNIST数据集的图像为28x28单通道,所以我们对LeNet-5网络结构进行了轻微的调整。下面是对修改后的完整代码,请参考。
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
torch.__version__
BATCH_SIZE=512 #大概需要2G的显存
EPOCHS=20 # 总共训练批次
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 让torch判断是否使用GPU,建议使用GPU环境,因为会快很多
#下载训练集
train_loader = torch.utils.data.DataLoader(
datasets.MNIST('data', train=True, download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])),
batch_size