CV-Implement【1】：Mnist

---

前言

Mnist包括6万张28x28的训练样本，1万张测试样本，可以说是计算机视觉领域的"Hello world"。

基于此，本文使用PyTorch复现Mnist，并就代码中的一些语法作展开。

---

1. 设置环境

import torch
import torchvision
from torch.backends import cudnn

import matplotlib.pyplot as plt

---

2. 准备数据集

2.1. 设置参数

# 定义了我们将在完整的训练数据集上循环多少次
n_epochs = 3
# 训练和测试批次的大小
batch_size_train = 64
batch_size_test = 1024
# optimizer的超参数
learning_rate = 0.01
momentum = 0.5
# 日志间隙，每隔 batch_size x log_interval 个batch才会输出一次loss
log_interval = 10
# 为了进行可重复的实验，我们必须为任何使用随机数生成的东西设置随机种子
random_seed = 1
torch.manual_seed(random_seed)
# 非确定性算法
torch.backends.cudnn.enabled = False
# 使用GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device

• batch
  在一个完整的训练集上训练一次，是一个epoch；而训练完一轮所有的batch是一个epoch
  • 每训练完一个batch，计算一次loss，进行一次梯度下降
  • shuffle：在训练每一轮epoch前，都会通过设置的shuffle参数随机打乱数据并分入不同的batch，即每一轮epoch中的batch的内容都不相同
• momentum
  可以用物理中的动量来类比，梯度在下降的时候，既要参考计算出来的梯度，也要参考前一次梯度下降的方向
  • 在一定程度上避免梯度下降到Local minimum就不再下降了
    
• cudnn
  cuDNN使用非确定性算法，使用时通常应该遵循以下准则：
  • 如果网络的输入数据维度或类型上变化不大，设置 torch.backends.cudnn.benchmark = true 可以增加运行效率；
  • 如果网络的输入数据在每次 iteration 都变化的话，会导致 cnDNN 每次都会去寻找一遍最优配置，这样反而会降低运行效率。

2.2. 加载数据集

train_loader = torch.utils.data.DataLoader(
  torchvision.datasets.MNIST('data/', train=True, download=True,
                             transform=torchvision.transforms.Compose([
                               torchvision.transforms.ToTensor(),
                               torchvision.transforms.Normalize(
                                 (0.1307,), (0.3081,)) # 用于Normalize()转换的值0.1307和0.3081是MNIST数据集的全局平均值和标准差
                             ])),
  batch_size=batch_size_train, shuffle=True)

test_loader = torch.utils.data.DataLoader(
  torchvision.datasets.MNIST('data/', train=False, download=True,
                             transform=torchvision.transforms.Compose([
                               torchvision.transforms.ToTensor(),
                               torchvision.transforms.Normalize(
                                 (0.1307,), (0.3081,)) # 用于Normalize()转换的值0.1307和0.3081是MNIST数据集的全局平均值和标准差
                             ])),
  batch_size=batch_size_test, shuffle=True)

2.3. 查看数据集的内容

examples = enumerate(test_loader)
batch_idx, (example_data, example_targets) = next(examples)

example_data.shape

test数据集的 batch_size = 1024，从这里我们可以发现，1个batch的测试数据是一个形状为 (1024, 1, 28, 28) 的张量，每个example都是大小为 28 x 28的灰度图

打印几个例子

fig = plt.figure()
for i in range(6):
  plt.subplot(2,3,i+1)
  plt.tight_layout()
  plt.imshow(example_data[i][0], cmap='gray', interpolation='none')
  plt.title("Ground Truth: {}".format(example_targets[i]))
  plt.xticks([])
  plt.yticks([])
fig

---

3. 搭建网络

3.1. 导入必要的库

导入一些子模块，以获得更多可读的代码

import torch.nn as nn
import torch.nn.functional as f
import torch.optim as optim

3.2. 网络模型

我们将使用两个二维卷积层，然后是两个全连接层。然后我们将选择 ReLU 作为激活函数。最后我们将使用两个 dropout 层作为正则化的手段。

class Net(nn.Module):
    def __init__(self):
        super().__init__() # python3.的写法
        # super(Net, self).__init__() - python2.的写法
        # 灰度图，通道数为1；10个卷积核，即输出的维度；卷积核的大小
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5) # (10, 24, 24)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5) # (20, 24, 24)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = f.relu(f.max_pool2d(self.conv1(x), 2)) #
        x = f.relu(f.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = f.relu(self.fc1(x))
        x = f.dropout(x, training=self.training)
        x = self.fc2(x)
        return f.log_softmax(x)

functional.max_pool2d & MaxPool2d

• functional.max_pool2d
  pytorch中的函数，可以直接调用

input = torch.randn(20, 16, 50, 32)  # 输入张量
f.max_pool2d(input, kernel_size=2, stride=1, padding=0)

• MaxPool2d
  pytorch中的类模块，先实例化，再调用其函数

m_1 = torch.nn.MaxPool2d(3, stride=2)  # 实例化
# 或者
m_2 = torch.nn.MaxPool2d((3, 2), stride=(2, 1))  # 实例化
input = torch.randn(20, 16, 50, 32)  # 输入张量
output = m_1(input) # 使用该类

初始化网络和优化器

network = Net()
optimizer = optim.SGD(network.parameters(), lr=learning_rate,
                      momentum=momentum)

---

4. 训练模型

4.1. 保存损失

我们还将通过一些打印输出来跟踪进度。为了在以后创建一个漂亮的训练曲线，我们还创建了两个列表来保存训练和测试损失。在X轴上，我们要显示网络在训练期间看到的训练实例的数量。

train_losses = []
train_counter = []
test_losses = []
test_counter = [i*len(train_loader.dataset) for i in range(n_epochs + 1)]

4.2. 训练模型

def train(epoch): # 在每个epoch上迭代一次所有的训练数据
    network.train() # 确保我们的网络处于训练模式
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad() # 手动将梯度设置为零
        output = network(data) # 产生网络的输出（前向传递）
        loss = f.nll_loss(output, target) # 计算损失
        loss.backward() # 反向传播
        optimizer.step() # 更新优化器
        if batch_idx % log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
            epoch, batch_idx * len(data), len(train_loader.dataset),
            100. * batch_idx / len(train_loader), loss.item()))
            train_losses.append(loss.item())
            train_counter.append(
            (batch_idx*64) + ((epoch-1)*len(train_loader.dataset)))

4.3. 测试模型

def test():
    network.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad(): # 使用上下文管理器no_grad()，避免在测试循环中更新模型参数
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = network(data)
            test_loss += f.nll_loss(output, target, size_average=False).item()
            pred = output.data.max(1, keepdim=True)[1]
            correct += pred.eq(target.data.view_as(pred)).sum()
    test_loss /= len(test_loader.dataset)
    test_losses.append(test_loss)
    print('\nTest set: Avg. loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))

4.4. 开始训练

test()
for epoch in range(1, n_epochs + 1):
    train(epoch)
    test()

torch.save(network.state_dict(), 'model/mnist/mnist_model.pth') # 只保存模型的参数，不保存模型
torch.save(optimizer.state_dict(), 'model/mnist/mnist_optimizer.pth') # 保存优化器的参数，如学习率等

---

5. 评估模型的性能

绘制训练曲线

fig = plt.figure()
plt.plot(train_counter, train_losses, color='blue')
plt.scatter(test_counter, test_losses, color='red')
plt.legend(['Train Loss', 'Test Loss'], loc='upper right')
plt.xlabel('number of training examples seen')
plt.ylabel('negative log likelihood loss')
fig

进行预测

fig = plt.figure()
for i in range(6):
    plt.subplot(2,3,i+1)
    plt.tight_layout()
    plt.imshow(example_data[i][0], cmap='gray', interpolation='none')
    plt.title("Prediction: {}".format(
        output.data.max(1, keepdim=True)[1][i].item()))
    plt.xticks([])
    plt.yticks([])
fig

---

6. 继续训练模型

实例化新的网络与新的优化器

continued_network = Net()
continued_optimizer = optim.SGD(network.parameters(), lr=learning_rate,
                                momentum=momentum)

加载第一次训练之后的模型和优化器的参数

network_state_dict = torch.load("model/mnist/mnist_model.pth")
continued_network.load_state_dict(network_state_dict)

optimizer_state_dict = torch.load("model/mnist/mnist_optimizer.pth")
continued_optimizer.load_state_dict(optimizer_state_dict)

继续训练

for i in range(4,9):
    test_counter.append(i*len(train_loader.dataset))
    train(i)
    test()

# 保存模型
torch.save(continued_network.state_dict(), 'data/result/mnist_model.pt') # 只保存模型的参数，不保存模型
torch.save(continued_optimizer.state_dict(), 'data/result/mnist_optimizer.pt') # 保存优化器的参数，如学习率等

打印新的损失函数图

fig = plt.figure()
plt.plot(train_counter, train_losses, color='blue')
plt.scatter(test_counter, test_losses, color='red')
plt.legend(['Train Loss', 'Test Loss'], loc='upper right')
plt.xlabel('number of training examples seen')
plt.ylabel('negative log likelihood loss')

---

总结

作为计算机视觉的"Hello World"，mnist无论从模型的深度还是代码的复杂程度都比较低，容易上手和理解。
同时，在这里尝试使用cuda对训练进行加速，取得了不错的成果。

参考资料