使用PyTorch实现简单的卷积神经网络识别MNIST手写数字

Pytorch 版本 0.4.1，系统：Ubuntu 16.04，Python版本：3.6
最近开始入门PyTorch，用了几天的时间学习，写了一个简单的手写数字识别练练手，参考了一个github的代码
现在做一个小小的总结和记录

导入模块及定义全局变量

程序使用到的模块有：

import torch
from torch import nn, optim # 网络及优化器
import torch.nn.functional as F # 网络中不需要更新参数的函数
from torch.autograd import Variable
from torchvision import datasets # 处理数据集
from torch.utils.data import DataLoader # 加载数据集
from torchvision import transforms 

batch_size = 128
learning_rate = 0.01 # 梯度下降的学习率
num_epoch = 10

加载数据集

pytorch需要先实例化数据集对象（MNIST是pytorch内置的对象，可以直接使用，如果是非内置的数据集需要自己编写数据集的类），再实例化DataLoader类，就可以用for循环提取数据样本了。

# 实例化MNIST数据集对象
train_data = datasets.MNIST('./dataset', train=True, transform=transforms.ToTensor(), download=True)
test_data = datasets.MNIST('./dataset', train=False, transform=transforms.ToTensor(), download=True)
# './dataset'是数据集存储的位置，相对于当前工程目录的路径，如果数据集已经存在则不会执行下载
# train_loader：以batch_size大小的样本组为单位的可迭代对象
train_loader = DataLoader(train_data, batch_size, shuffle=True)
test_loader = DataLoader(test_data)

关于pytorch内置的数据集和用法参见：PyTorch中文文档

定义模型

模型是程序的核心部分，但是这里重点讨论PyTorch的基础使用，因此在模型的优化上不作过多介绍，只使用最基础的CNN和全连接层。结构如下：

卷积 > 卷积 > 池化 > 全连接 > 全连接 > 全连接

class CNN(nn.Module):
    def __init__(self, in_dim, out_dim):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(in_dim, 6, 3, stride=1, padding=1)
        self.batch_norm1 = nn.BatchNorm2d(6)
        self.relu = nn.ReLU(True) # 激活函数使用ReLU
        self.conv2 = nn.Conv2d(6, 16, 5, stride=1, padding=0)
        self.pool = nn.MaxPool2d(2, 2)
        self.batch_norm2 = nn.BatchNorm2d(16)

        self.fc1 = nn.Linear(400, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, out_dim)

    def forward(self, x):
        x = self.batch_norm1(self.conv1(x))
        x = self.relu(x) # 这里也可以直接写作x = F.relu(x)，现在我还不知道这两种方法有什么区别
        x = self.pool(x)
        x = self.batch_norm2(self.conv2(x))
        x = self.relu(x)
        x = self.pool(x)
        x = x.view(x.size(0), -1) # reshape操作，将x展开
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

    def print_model_name(self):
        print("Model Name: CNN")

还可以使用nn.Sequential函数构建：

class Cnn(nn.Module):
    def __init__(self, in_dim, n_class):
        super(Cnn, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_dim, 6, 3, stride=1, padding=1),
            nn.ReLU(True),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(6, 16, 5, stride=1, padding=0),
            nn.ReLU(True),
            nn.MaxPool2d(2, 2))

        self.fc = nn.Sequential(
            nn.Linear(400, 120), nn.Linear(120, 84), nn.Linear(84, n_class))

    def forward(self, x):
        # print(x.size()) torch.Size([128, 1, 28, 28])
        out = self.conv(x)
        out = out.view(out.size(0), -1)
        # print(out.size()) = torch.Size([128, 400])
        out = self.fc(out)
        # print(out.size()) torch.Size([128, 10])
        return out

以上代码来自开头的参考代码，网络有一点不同，影响不大。可以看出，nn.Sequential()函数可以直接将每一层连接起来， 简化代码。

训练

训练过程比较简单。

• 检测GPU是否可用

isGPU = torch.cuda.is_available()
print(isGPU)

• 实例化模型

model = CNN(1, 10)
if isGPU:
    model = model.cuda()

• 设置损失函数，优化器

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate)

这里使用的是交叉熵函数和随机梯度下降。

• 训练过程

直接上代码

for epoch in range(num_epoch):
    running_acc = 0.0
    running_loss = 0.0
    for i, data in enumerate(train_loader, 1): # train_loader：以batch_size大小的样本组为单位的可迭代对象
        img, label = data
        img = Variable(img)
        label = Variable(label)
        if isGPU:
            img = img.cuda()
            label = label.cuda()
        # forward
        out = model(img) # 得到网络的输出
        loss = criterion(out, label)
        # backward
        optimizer.zero_grad() # 梯度值归0，这一步不可缺少
        loss.backward()
        optimizer.step()
	
        _, pred = torch.max(out, dim=1)  # 按维度dim 返回最大值及索引
        running_loss += loss.item()*label.size(0)
        current_num = (pred == label).sum() # 每个batch预测正确的个数，类型为variable
        acc = (pred == label).float().mean()   # 每个batch预测正确率，类型为variable
        running_acc += current_num.item()

        if i % 100 == 0: # 每100个batch打印一次结果
            print("epoch: {}/{}, loss: {:.6f}, running_acc: {:.6f}"
                  .format(epoch+1, num_epoch, loss.item(), acc.item()))
        # loss是Variable类型，调用item()方法得到数值
    print("epoch: {}, loss: {:.6f}, accuracy: {:.6f}".format(epoch+1, running_loss, running_acc/len(train_data)))

• 测试

需要首先进入测试模式，后面的过程和训练差不多。

model.eval() # 测试模式
current_num = 0
for i , data in enumerate(test_loader, 1):
    img, label = data
    if isGPU:
        img = img.cuda()
        label = label.cuda()
    with torch.no_grad(): # 测试的时候不需要计算梯度
        img = Variable(img)
        label = Variable(label)
    out = model(img)
    _, pred = torch.max(out, 1)
    current_num += (pred == label).sum().item()

print("Test result: accuracy: {:.6f}".format(float(current_num/len(test_data))))

• 最后，保存模型

torch.save(model.state_dict(), './cnn.pth') # 保存模型

完整代码：https://github.com/chenyr0021/basicCNN-pytorch

遇到的问题

• 运行问题

RuntimeError: “log_softmax_lastdim_kernel_impl” not implemented for ‘torch.LongTensor’

后来发现是criterion函数执行时会自动将label参数进行one-hot运算，所以out 和label的维度是不一样的。

print(out.size()) # torch.Size([128, 10])
print(label.size()) # torch.Size([128])

参考：https://stackoverflow.com/questions/51818225/pytorch-runtimeerror-host-softmax-not-implemented-for-torch-cuda-longtensor

• 监视GPU运行情况

watch nvidia-smi

watch 命令的使用：
watch [options] command
最常用的参数是 -n， 后面指定是每多少秒来执行一次命令。
监视显存
我们设置为每 10s 显示一次显存的情况
watch -n 10 nvidia-smi
参考：https://blog.csdn.net/wuguangbin1230/article/details/72886903