【Pytorch】4 训练一个图像分类器（Pytorch官方手册CIFAR10案例）

构建图像分类网络（Pytorch官方手册CIFAR10案例）

引言

对于视觉，我们已经创建了一个叫做 totchvision的包，该包含有支持加载类似Imagenet，CIFAR10，MNIST等公共数据集的数据加载模块 torchvision.datasets和支持加载图像数据数据转换模块 torch.utils.data.DataLoader。

CIFAR10数据集，它包含十个类别：‘airplane’, ‘automobile’, ‘bird’, ‘cat’, ‘deer’, ‘dog’, ‘frog’, ‘horse’, ‘ship’, ‘truck’。CIFAR-10 中的图像尺寸为3 * 32 * 32，也就是RGB的3层颜色通道，每层通道内的尺寸为32*32。

---

训练一个图像分类器主要有以下步骤

• 使用torchvision加载并归一化CIFAR10的训练和测试数据集
• 定义一个卷积神经网络
• 定义一个损失函数
• 在训练样本数据上训练网络
• 在测试样本数据上测试网络

---

第一部分：加载数据集

完整代码：

import torch
import torchvision
import torchvision.transforms as transforms

transform = transforms.Compose(
    [
        transforms.ToTensor(),
        # 归一化 平均值0.5 标准差0.5 后续为什么反归一化为 img = img / 2 + 0.5 原因在这里
        # 这里应该是 batch Normalization 为了使训练更容易 使用归一化
        #                        x - mean
        # Batch Normalization = ----------
        #                          std
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ]
)

transet = torchvision.datasets.CIFAR10(root='./data/train', train=True, download=True, transform=transform)
testset = torchvision.datasets.CIFAR10(root='./data/test', train=False, download=True, transform=transform)

trainloader = torch.utils.data.DataLoader(transet, batch_size=4, shuffle=True)
testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=False)

classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

---

注释一：
transforms模块主要包括对图像的一些操作：

详情可以参考这位大佬的博客：torchvision库及其常用函数

transform = transforms.Compose(
    [
        transforms.ToTensor(),
        # 归一化 平均值0.5 标准差0.5 后续为什么反归一化为 img = img / 2 + 0.5 原因在这里
        # 这里应该是 batch Normalization 为了使训练更容易 使用归一化
        #                        x - mean
        # Batch Normalization = ----------
        #                          std
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ]
)

注释二：加载数据集操作

torchvision.datasets.*用于加载常用数据集，Pytorch目前支持31种数据集：

- Caltech、CelebA、CIFAR、Cityscapes、COCO、EMNIST
- FakeData、Fashion-MNIST、Flickr、HMDB51、ImageNet、iNaturalist
- Kinetics-400、KITTI、KMNIST、LFW、LSUN
- MNIST、Omniglot、PhotoTour、Places365、QMNIST
- SBD、SBU、SEMEION、STL10、SVHN
- UCF101、USPS、VOC、WIDERFace、Base classes for custom datasets

torchvision.datasets.CIFAR10参数：

• root：数据集下载到的根目录，数据集所在的根目录，这里是cifar-10-batches-py的根目录
• train：True代表加载训练集，False代表加载测试集
• download：True从互联网上下载数据，并将其放在root目录下，如果数据集已经下载，什么也不干。（这里下载可能速度慢，需要魔法）链接：CIFAR
  提取码：8lwp
• trasnform：设置一组对图像进行处理的操作，这一组操作由Compose组成,这一组compose 的顺序还很重要。

torchvision.data.Dataloader参数：

• trainset：将刚刚生成的trainset列表传入torchvision.data.Dataloader的第一个参数。
• batch size：设定图像批次大小，就是一次处理几张图片
• shuffle：True代表每一个epoch过程中会打乱数据顺序，重新随机选择，一般都会设置为True
• num worker：导入数据时用几个线程操作，windows一般设置为0，我用的windows设置为2时程序报错，Linux可以设成2.

transet = torchvision.datasets.CIFAR10(root='./data/train', train=True, download=True, transform=transform)
testset = torchvision.datasets.CIFAR10(root='./data/test', train=False, download=True, transform=transform)

trainloader = torch.utils.data.DataLoader(transet, batch_size=4, shuffle=True)
testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=False)

---

第二部分：训练数据

完整代码：

class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()

        self.conv1 = nn.Conv2d(3, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)

        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, self.num_flat_features(x))
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

    def num_flat_features(self, x):
        size = x.size()[1:]
        num_features = 1
        for s in size:
            num_features *= s
        return num_features

device = torch.device('cpu:0')
# 第一步：创建网络
net = Net()
# net.to(device)
# 第二部：创建损失函数 这里用的交叉熵
criterion = nn.CrossEntropyLoss()
# criterion.to(device)
# 第三步：创建权重更新规则 这里用的SGD
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
# 第四步：循环迭代数据 在数据迭代器上循环传给网络和优化器输入
since = time.time()
for epoch in range(2):
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data
        # inputs = inputs.to(device)
        # labels = labels.to(device)

        # 核心代码
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        if i % 2000 == 1999:
            print('[{0},{1}] loss:{2}'.format(epoch + 1, i + 1, running_loss / 2000))
            running_loss = 0.0
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(
    time_elapsed // 60, time_elapsed % 60))
print('Finished Training')

---

注释一：创建网络、损失函数、设置权重更新规则

# 第一步：创建网络
net = Net()
# 第二部：创建损失函数 这里用的交叉熵
criterion = nn.CrossEntropyLoss()
# 第三步：创建权重更新规则 这里用的SGD
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

注释二：循环迭代数据

# 迭代2次epoch

以上都是必要步骤，核心代码。

注释三：可能会经常用到的代码

1. 更新权重核心代码

# 核心代码
optimizer.zero_grad()
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()

2. 计算训练时间核心代码

import time

since = time.time()

for epoch in range(2):
    ···
    
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(
    time_elapsed // 60, time_elapsed % 60))

3. 显示图像

def imshow(img):
    # 反归一化
    img = img / 2 + 0.5
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()

    ···
    
# 方法一：官方实例
dataiter = iter(trainloader)
images, labels = dataiter.next()
imshow(torchvision.utils.make_grid(images))
print(' '.join('%5s' % classes[labels[j]] for j in range(4)))

# 方法二：网上看到的
for image, label in trainloader:
    imshow(image[0])

注释四：关于GPU加速

关于如何使用GPU加速，主要对Net``````criterion``````image``````labels进行to(device)操作即可。

也可以使用.cuda()，例如net.cuda()，但看网上博客说不推荐使用，尽量使用to(device)。

我用的GPU是RTX3060，目前使用GPU加速与不使用GPU加速我测试的训练时间基本一样，稳定在52s，没有体现出来GPU加速的效果。我猜测可能还是因为网络比较简单或者数据集数量比较小。或者还是我的代码有错误，如果有错误，希望大佬们可以评论指出。

另外一个是我用RTX3060遇到了显卡不兼容的问题，很有可能是Pytorch版本太低了，我用的CUDA是11.3，但在之前使用的Pytorch是1.9.0版本，升级到Pytorch 1.10.0就顺利解决问题了。
参考博客：显卡不兼容问题

我直接搬运过来了，如果你也是RTX3060，直接在Pytorch所在环境下运行下面代码即可：

# 使用conda卸载Pytorch 
conda uninstall pytorch
conda uninstall libtorch

# 使用pip卸载Pytorch
pip uninstall torch

# 安装Pytorch 1.10.0 + gpu版本 需要CUDA 11.3版本
pip3 --default-timeout=60000 install torch==1.10.0+cu113 torchvision==0.11.1+cu113 torchaudio===0.10.0+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html -i https://pypi.tuna.tsinghua.edu.cn/simple --trusted-host pypi.tuna.tsinghua.edu.cn

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

net = Net()
net.to(device)

criterion = nn.CrossEntropyLoss()
criterion.to(device)

··

for epoch in range(2):
    ···
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data
        inputs = inputs.to(device)
        labels = labels.to(device)
    ···

---

第三部分：程序完整代码

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable

import torchvision
import torchvision.transforms as transforms

import matplotlib.pyplot as plt
import numpy as np

import torch.optim as optim

import time


class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()

        self.conv1 = nn.Conv2d(3, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)

        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, self.num_flat_features(x))
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

    def num_flat_features(self, x):
        size = x.size()[1:]
        num_features = 1
        for s in size:
            num_features *= s
        return num_features


def imshow(img):
    # 反归一化
    img = img / 2 + 0.5
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()


transform = transforms.Compose(
    [
        transforms.ToTensor(),
        # 归一化 平均值0.5 标准差0.5 后续为什么反归一化为 img = img / 2 + 0.5 原因在这里
        # 这里应该是 batch Normalization 为了使训练更容易 使用归一化
        #                        x - mean
        # Batch Normalization = ----------
        #                          std
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ]
)

transet = torchvision.datasets.CIFAR10(root='./data/train', train=True, download=True, transform=transform)
testset = torchvision.datasets.CIFAR10(root='./data/test', train=False, download=True, transform=transform)

trainloader = torch.utils.data.DataLoader(transet, batch_size=4, shuffle=True)
testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=False)

classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

# 展示一些图片
# # get some random training images
# dataiter = iter(trainloader)
# images, labels = dataiter.next()
#
# # show images
# imshow(torchvision.utils.make_grid(images))
# # print labels
# print(' '.join('%5s' % classes[labels[j]] for j in range(4)))


# for image, label in trainloader:
#     imshow(image[0])

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# 第一步：创建网络
net = Net()
# net.to(device)
# 第二部：创建损失函数 这里用的交叉熵
criterion = nn.CrossEntropyLoss()
# criterion.to(device)
# 第三步：创建权重更新规则 这里用的SGD
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
# 第四步：循环迭代数据 在数据迭代器上循环传给网络和优化器输入
since = time.time()
for epoch in range(2):
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data
        # inputs = inputs.to(device)
        # labels = labels.to(device)

        # 核心代码
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        if i % 2000 == 1999:
            print('[{0},{1}] loss:{2}'.format(epoch + 1, i + 1, running_loss / 2000))
            running_loss = 0.0
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(
    time_elapsed // 60, time_elapsed % 60))
print('Finished Training')