手写数字识别 卷积神经网络 Pytorch框架实现

MNIST 手写数字识别 卷积神经网络 Pytorch框架

谨此纪念刚入门的我在卷积神经网络上面的摸爬滚打

说明

下面代码是使用pytorch来实现的LeNet，可以正常运行测试，自己添加了一些注释，方便查看。

代码实现

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms

# Device configuration
#这里是个python的三元表达式，如果cuda存在的话，divice='cuda:0'，否者就是'cpu'
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

# Hyper parameters
num_epochs = 5  #全部训练集使用的次数
num_classes = 10 #全连接层输出的结果种类
batch_size = 100  #批处理的图片的个数
learning_rate = 0.001 #学习率，在梯度下降法里面的系数

# MNIST dataset
#下载训练数据集，位置放在本文件的父文件夹下的data文件夹里面，数据需要转换格式为Tensor
train_dataset = torchvision.datasets.FashionMNIST(root='../data/',
                                           train=True,
                                           transform=transforms.ToTensor(),
                                           download=True)
#下载测试集，位置放在放在本文件的父文件夹下的data文件夹里面，数据需要转换为Tensor格式
test_dataset = torchvision.datasets.FashionMNIST(root='../data/',
                                          train=False,
                                          transform=transforms.ToTensor())

# Data loader
#这里的shuffle(bool, optional):在每个epoch开始的时候，对数据进行重新打乱，就是重新分组
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                           batch_size=batch_size,
                                           shuffle=True)
#对于测试集来说不需要进行从新分组(这里好像不是必须的，也可以试试每次测试分组，有意义吗？)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                          batch_size=batch_size,
                                          shuffle=False)

# Convolutional neural network (two convolutional layers)
#定义一个卷积类，这里需要继承nn.Module，它是专门为神经网络设计的模块化接口
class ConvNet(nn.Module):
    def __init__(self, num_classes=10):
        #调用父类的初始化函数
        super(ConvNet, self).__init__()
        #一个有序的容器，神经网络模块将按照在传入构造器的顺序依次被添加到计算图中执行，同时以神经网络模块为元素的有序字典也可以作为传入参数。
        self.layer1 = nn.Sequential(
            #二维卷积层，输入通道数1，输出通道数16(相当于有16个filter，也就是16个卷积核)，卷积核大小为5*5，步长为1，零填充2圈
            #经过计算，可以得到卷积输出的图像的大小和输入的图像大小是等大小的，但是深度不一样，为28*28*16（16为深度），因为这里的padding抵消了卷积的缩小
            nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2),
            #BatchNorm2d是卷积网络中防止梯度消失或爆炸的函数，参数是卷积的输出通道数
            nn.BatchNorm2d(16),
            #激活函数
            nn.ReLU(),
            #二维最大池化，核的大小为2，步长为2
            #这样输出的图片大小就是14*14*16（16为深度）
            nn.MaxPool2d(kernel_size=2, stride=2))
        #两层的卷积网络，具体含义和上面相同
        self.layer2 = nn.Sequential(
            #这里大小也没有变化，输出依然和输出的大小相同，深度为32，所以图像为14*14*32
            #但是这里的卷积核的数量是32，和输出通道数相同。
            nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            #下面经过池化后输出就会变成7*7*32
            nn.MaxPool2d(kernel_size=2, stride=2))
        #对输入数据做线性变换，第一个参数是每个输入样本的大小：7*7*32；第二个参数是输出样本的大小，这里是10，正好代表10个数，相当于类别
        #第三个参数为bias（偏差），默认为True。如果为False，那么这层将不会学习偏置。
        self.fc = nn.Linear(7*7*32, num_classes)

    #定义了每次执行的 计算步骤。 在所有的子类中都需要重写forward函数。
    #
    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = out.reshape(out.size(0), -1)
        out = self.fc(out)
        return out
model = ConvNet(num_classes).to(device)

# Loss and optimizer
#损失函数，
criterion = nn.CrossEntropyLoss()
#优化函数
#params (iterable)第一个参数：待优化参数的iterable或者是定义了参数组的dict
#lr (float, 可选) – 学习率（默认：1e-3）同样也称为学习率或步长因子，它控制了权重的更新比率（如 0.001）。
#较大的值（如 0.3）在学习率更新前会有更快的初始学习，而较小的值（如 1.0E-5）会令训练收敛到更好的性能。
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

# Train the model
#total_step是每一轮的测试次数，这里就是60000/100=600次
total_step = len(train_loader)
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        images = images.to(device)
        labels = labels.to(device)

        # Forward pass
        #model(images)等价于module.forward(images)
        outputs = model(images)
        #根据输出的结果和标签对比，计算loss
        loss = criterion(outputs, labels)

        # Backward and optimize
        #根据pytorch中的backward()函数的计算，当网络参量进行反馈时，梯度是被积累的而不是被替换掉；
        #但是在每一个batch时毫无疑问并不需要将两个batch的梯度混合起来累积，因此这里就需要每个batch设置一遍zero_grad 了
        #将梯度初始化为零
        optimizer.zero_grad()
        #这里是使用反向传播计算梯度值
        loss.backward()
        #在scheduler的step_size表示scheduler.step()每调用step_size次，对应的学习率就会按照策略调整一次
        optimizer.step()

        if (i+1) % 100 == 0:
            print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'
                   .format(epoch+1, num_epochs, i+1, total_step, loss.item()))

# Test the model
model.eval()  # eval mode (batchnorm uses moving mean/variance instead of mini-batch mean/variance)
with torch.no_grad():
    correct = 0
    total = 0
    for images, labels in test_loader:
        images = images.to(device)
        labels = labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    print('Test Accuracy of the model on the 10000 test images: {} %'.format(100 * correct / total))

# Save the model checkpoint
torch.save(model.state_dict(), 'model.ckpt')