【Pytorch教程】08-如何使用PyTorch训练简单CIFAR10图片分类器(保姆级)

1. 背景

1.1 各领域常用库

在深度学习的实战中，通常要处理各种类型的数据，例如图片、文本、声音和视频。我们可以使用Python标准库来将这些数据加载到 Numpy 数组中，然后通过[3.1.2 通过Numpy arrays创建张量](# 3.1.2 通过Numpy arrays创建张量)学习的内容把 Numpy 数组转换为PyTorch能处理的张量 torch.Tensor 。

下面总结了在图像、音频和文本领域常用的Python库：

领域	常用的Python库
图像	Pillow, OpenCV
音频	SciPy, Librosa
文本	原生Python, Cython, NLTK, SpaCy

非常幸运的是，PyTorch框架为专为计算机视觉领域设计了一个非常好用且方便的库 torchvision 。 torchvision 里面封装了诸如 ImageNet、CIFAR10、MNIST 等各种常见数据集的数据集加载模块 (data loaders) ，并且还封装了用于图像的数据转换模块。即 torchvision.datasets 和 torch.utils.data.DataLoader 。

这些PyTorch模块提供了巨大的便利，极大地避免了编写冗余的代码。

在本次实战教程中，我们将以 CIFAR10 数据集为例，带领大家从数据集加载、搭建神经网络、训练神经网络模型和参数模型准确率，让大家全流程体验PyTorch给深度学习带来的便捷。

1.2 CIFAR10简介

CIFAR10数据集一共有10个类别，分别是 ‘airplane’, ‘automobile’, ‘bird’, ‘cat’, ‘deer’, ‘dog’, ‘frog’, ‘horse’, ‘ship’, ‘truck’ 。

每张图片的尺寸为 $3\times 32\times 32$ ，其中，$3$ 代表是三通道的彩色图片；$32$ 代表图片的高 (height) 和宽 (width) 都是32像素。

本次的实战将按下列步骤展开：

1. 使用PyTorch的 torchvision 库加载并正则化CIFAR10的训练集和测试集；
2. 搭建一个卷积神经网络；
3. 定义损失函数；
4. 在训练集上训练神经网络模型；
5. 在测试集上测试模型的识别准确率。

2. 数据集

2.1 加载并正则化CIFAR10数据集

借助 torchvision 库，我们很容易就能加载 CIFAR10 数据集：

import torch
import torchvision
import torch.utils.data
import torchvision.transforms as transforms

【注意】

使用 torch.utils.data.DataLoader() 时，最好在开头加一句 import torch.utils.data 。否则PyCharm可能会识别不出 data 发出警告。

【注意】

• torchvision 库输出的数据集都是像素值范围在 [0, 1] 的 PIL Image 格式图片。我们将它们转换为张量归一化的范围： [-1, 1] 。
• 如果你在Windows系统上运行或调试Debug，可能会出现 BrokenPipeError 或者 Connected 无反应的报错。可以试试将 torch.utils.data.DataLoader() 的 num_worker 设置成 0 。

# 封装一组转换函数对象作为转换器
transform = transforms.Compose(  # Compose是transforms的组合类
    [transforms.ToTensor(),  # ToTensor()类把PIL Image格式的图片和Numpy数组转换成张量
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]  # 用均值和标准差归一化张量图像
)

# 声明批量大小，一批4张图片
batch_size = 4

# 实例化训练集
train_set = torchvision.datasets.CIFAR10(root='./data', train=True,
                                         transform=transform, download=True)

# 实例化训练集加载器
train_loader = torch.utils.data.DataLoader(train_set, batch_size=batch_size,
                                           shuffle=True, num_workers=12)

# 实例化测试集
test_set = torchvision.datasets.CIFAR10(root='./data', train=False,
                                        transform=transform, download=True)

# 实例化测试集加载器
test_loader = torch.utils.data.DataLoader(test_set, batch_size=batch_size,
                                          shuffle=True, num_workers=12)

# CIFAR10数据集所有类别名称
classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

输出：

Files already downloaded and verified
Files already downloaded and verified

2.2 训练集可视化

前面提到，torchvision 库输出的数据集都是像素值范围在 [0, 1] 的 PIL Image 格式图片。如果想要输出CIFAR10训练集的原始图片，需要如下代码所示：

import matplotlib.pyplot as plt
import numpy as np

def img_show(img):
    img = img / 2 + 0.5  # 反正则化
    npimg = img.numpy()  # 转换成Numpy数组
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()


dataiter = iter(train_loader)       # 训练集图片的迭代器
images, labels = dataiter.next()    # 获取每个图片和标签
img_show(torchvision.utils.make_grid(images))   # 显示图片
print(" ".join(f'{classes[labels[j]]:5s}' for j in range(batch_size)))  # 打印标签

输出图片：

输出标签：

truck  dog   dog   bird 

3. 定义卷积神经网络

• 为了简单起见，定义一个简单的卷积神经网络：

  

  import torch.nn as nn
  import torch.nn.functional as F
  
  # 定义网络模型，继承自torch.nn.Module类
  class Model(nn.Module):
      # 构造器
      def __init__(self):
          super().__init__()  # 初始化父类的属性
          # Model类的属性
          self.conv1 = nn.Conv2d(3, 6, 5)  # 卷积层1
          self.pool = nn.MaxPool2d(2, 2)  # 最大池化层
          self.conv2 = nn.Conv2d(6, 16, 5)  # 卷积层2
          self.fc1 = nn.Linear(16 * 5 * 5, 120)  # 全连接层1
          self.fc2 = nn.Linear(120, 84)  # 全连接层2
          self.fc3 = nn.Linear(84, 10)  # 全连接层3
  
      # 前向传播方法
      def forward(self, x):
          x = self.pool(F.relu(self.conv1(x)))
          x = self.pool(F.relu(self.conv2(x)))
          x = torch.flatten(x, 1)  # 把x沿着水平方向展开
          x = F.relu(self.fc1(x))
          x = F.relu(self.fc2(x))
          x = self.fc3(x)
          return x
  
  
  # 实例化模型对象
  model = Model()
  

4. 定义损失函数和优化器

• 这里我们使用分类交叉熵损失函数，优化器使用带动量的随机梯度下降 (SGD) 。

  import torch.optim as optim
  
  # 定义损失函数
  loss = nn.CrossEntropyLoss()
  
  # 定义优化器
  optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
  

5. 训练模型

• 这是本节的精彩内容。我们只需简单地循环我们的数据遍历器，然后将训练集图片输入到模型中，然后它就会自动优化更新参数。

  # 开始训练
  for epoch in range(4):  # 训练4个epoch
  
      running_loss = 0.0  # 损失函数记录
      for i, data in enumerate(train_loader, 0):
          # 获取模型输入；data是由[inputs, labels]组成的列表
          inputs, labels = data
  
          # 把参数的梯度清零
          optimizer.zero_grad()
  
          # 前向传播+反向传播+更新权重
          outputs = model(inputs)
          loss = criterion(outputs, labels)
          loss.backward()
          optimizer.step()
  
          # 打印统计数据
          running_loss += loss.item()
          if i % 2000 == 1999:  # 每2000个mini-batch打印一次统计信息
              print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}')
              running_loss = 0.0  # 把损失函数记录清零
  
  print("训练结束")
  

6. 保存模型参数

• 马上保存我们训练得到的模型参数：

  # 保存模型参数
  PATH = './output/cifar_model.pth'  # 指定模型参数保存路径
  torch.save(model.state_dict(), PATH)  # 保存模型参数
  

• 接下来就可以开始训练了：

  [1,  2000] loss: 2.257
  [1,  4000] loss: 1.927
  [1,  6000] loss: 1.718
  [1,  8000] loss: 1.599
  [1, 10000] loss: 1.544
  [1, 12000] loss: 1.473
  [2,  2000] loss: 1.429
  [2,  4000] loss: 1.382
  [2,  6000] loss: 1.345
  [2,  8000] loss: 1.346
  [2, 10000] loss: 1.328
  [2, 12000] loss: 1.332
  [3,  2000] loss: 1.253
  [3,  4000] loss: 1.247
  [3,  6000] loss: 1.233
  [3,  8000] loss: 1.192
  [3, 10000] loss: 1.223
  [3, 12000] loss: 1.202
  [4,  2000] loss: 1.115
  [4,  4000] loss: 1.147
  [4,  6000] loss: 1.139
  [4,  8000] loss: 1.139
  [4, 10000] loss: 1.141
  [4, 12000] loss: 1.133
  训练结束
  

7. 测试

• 经过上面的 4 epoch 的训练，我们相当于把所有训练集图片学习了四次。接下来我们需要测试模型训练得怎么样。
• 我们会把测试集图片输入到训练好的模型中，然后模型就会输出类别的预测结果，把这个预测类别与图片的实际类别比较。如果模型预测正确，就把图片添加到正确预测的列表中。

7.1 测试集可视化

• 在此之前，我们可以随机显示一些测试集图片来查看：

  # 随机选取测试集图片
  dataiter = iter(test_loader)
  images, labels = dataiter.next()
  
  # 打印输出测试集图片
  img_show(torchvision.utils.make_grid(images))
  print('GroundTruth: ', ' '.join(f'{classes[labels[j]]:5s}' for j in range(4)))
  

  

  GroundTruth:  bird  bird  cat   car
  

7.2 加载模型参数

• 接下来，我们加载刚刚训练好的模型参数：

  # 实例化模型对象
  model = Model().to('cuda')
  
  # 加载模型
  model.load_state_dict(torch.load(PATH))
  

• 我们来看看卷积神经网络觉得测试集图片是什么类别叭：

  # 模型预测
  outputs = model(images)
  

• 预测输出的 10 个类别的概率。某个类别的概率值越高，说明模型觉得这个图片更有可能是属于这个类别。因此，我们直接获取最高概率值元素对应的索引：

  # 选择置信度最高的类别作为预测类别
  _, predicted = torch.max(outputs, 1)
  
  print('Predicted: ', ' '.join(f'{classes[predicted[j]]:5s}' for j in range(4)))
  

• 总体的测试过程代码如下：

  PATH = './outputs/cifar_model.pth'  # 指定模型参数保存路径
  # 加载模型参数
  model.load_state_dict(torch.load(PATH))
  correct = 0  # 预测正确的数量
  total = 0	 # 测试集的总数
  
  # 由于我们不是训练，我们不需要计算输出的梯度
  with torch.no_grad():
      for data in test_loader:
          images, labels = data
          images = images.to('cuda')
          labels = labels.to('cuda')
  
          # 模型输出预测结果
          outputs = model(images)
          
          # 选择置信度最高的类别作为预测类别
          _, predicted = torch.max(outputs.data, 1)
          total += labels.size(0)
          correct += (predicted == labels).sum().item()
  
  # 打印准确率
  print(f'Accuracy of the network on the 10000 test images: {100 * correct // total}%')
  

  输出：

  Accuracy of the network on the 10000 test images: 57%
  

  准确率为 57 % 。这比从 10 个类别盲猜一个的概率 (10%) 要大，说明模型确实是学习到一些东西的。

8. 多卡训练

• 如果你的机器有多块 GPU ，想在多卡机器上训练，欢迎参考我的这篇博文：【PyTorch教程】07-PyTorch如何使用多块GPU训练神经网络模型 。