PyTorch教程：训练分类器

June 09, 2020 (最后修改: August 21, 2021)

本文翻译自Pytorch官方文档《DEEP LEARNING WITH PYTORCH: A 60 MINUTE BLITZ》

就是这个。您已经了解了如何定义神经网络，计算损耗并更新网络的权重。

现在您可能在想，

数据呢？

通常，当您必须处理图像，文本，音频或视频数据时，可以使用标准 python 包将数据加载到 numpy 数组中。然后，您可以将此数组转换为 torch.*Tensor。

对于图像，有用的软件包如 Pillow，OpenCV
对于音频，请使用 scipy 和 librosa 等软件包
对于文本，基于原始 Python 或 Cython 加载，或者 NLTK 和 SpaCy 很有用

专门针对视觉，我们创建了一个名为 torchvision 的程序包，其中包含用于常见数据集（例如 Imagenet，CIFAR10，MNIST等）的数据加载器，以及用于图像（即 torchvision.datasets 和 torch.utils.data.DataLoader）的数据转换器。

这提供了极大的便利，并且避免了编写样板代码。

在本教程中，我们将使用 CIFAR10 数据集，包含以下类别： ‘airplane’, ‘automobile’, ‘bird’, ‘cat’, ‘deer’, ‘dog’, ‘frog’, ‘horse’, ‘ship’, ‘truck’。 CIFAR-10 中的图像尺寸为 3x32x32，即尺寸为 32x32 像素的 3 通道彩色图像。

训练图片分类器

我们将按顺序执行以下步骤：

使用 torchvision 加载和标准化 CIFAR10 训练和测试数据集
定义卷积神经网络
定义损失函数
根据训练数据训练网络
在测试数据上测试网络

1. 加载和标准化 CIFAR10

使用 torchvision，加载 CIFAR10 非常容易。

import torch
import torchvision
import torchvision.transforms as transforms

torchvision 数据集的输出是 [0，1] 范围的 PILImage 图像。我们将它们转换为归一化范围 [-1，1] 的 Tensor。

注意：

如果在 Windows 上运行并且出现 BrokenPipeError，请尝试将 torch.utils.data.DataLoader() 的 num_worker 设置为 0。

transform = transforms.Compose(
    [
        transforms.ToTensor(),
        transforms.Normalize(
            (0.5, 0.5, 0.5),
            (0.5, 0.5, 0.5),
        )
    ]
)

trainset = torchvision.datasets.CIFAR10(
    root="./data",
    train=True,
    download=True,
    transform=transform,
)
trainloader = torch.utils.data.DataLoader(
    trainset,
    batch_size=4,
    shuffle=True,
    num_workers=2,
)

testset = torchvision.datasets.CIFAR10(
    root="./data",
    train=False,
    download=True,
    transform=transform,
)
testloader = torch.utils.data.DataLoader(
    testset,
    batch_size=4,
    shuffle=False,
    num_workers=2,
)

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

让我们展示一些训练图像，很有趣。

import matplotlib.pyplot as plt
import numpy as np

# 显示图片的函数
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(f"{classes[labels[j]]:5s}" for j in range(4)))

2. 定义卷积神经网络

之前从“神经网络”部分复制神经网络，然后对其进行修改以获取 3 通道图像（而不是定义的 1 通道图像）。

import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        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 = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


net = Net()

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

让我们使用分类交叉熵损失和带有动量的 SGD。

import torch.optim as optim

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(
    net.parameters(), 
    lr=0.001, 
    momentum=0.9
)

4. 训练网络

这是事情开始变得有趣的时候。我们只需要遍历数据迭代器，然后将输入馈送到网络并进行优化。

for epoch in range(2):  # 多次遍历数据集
    
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        # 获取输入；data 是 [inputs, labels] 的列表
        inputs, labels = data
        
        # 参数梯度归零
        optimizer.zero_grad()
        
        # forward + backward + optimize
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        # 打印统计量
        running_loss += loss.item()
        if i % 2000 == 1999:  # 每 2000 个输出一次
            print(f"[{epoch + 1}, {i + 1:5d}] loss: {running_loss/2000:.3f}")
            running_loss = 0.0

print("Finished Training")


[1,  2000] loss: 2.219
[1,  4000] loss: 1.855
[1,  6000] loss: 1.683
[1,  8000] loss: 1.618
[1, 10000] loss: 1.547
[1, 12000] loss: 1.498
[2,  2000] loss: 1.445
[2,  4000] loss: 1.402
[2,  6000] loss: 1.373
[2,  8000] loss: 1.344
[2, 10000] loss: 1.329
[2, 12000] loss: 1.294
Finished Training

让我们快速保存我们训练过的模型：

path = "./cifar_net.pth"
torch.save(net.state_dict(), path)

有关保存PyTorch模型的更多详细信息，请参见下面的文档。

https://pytorch.org/docs/stable/notes/serialization.html

5. 在测试数据上测试网络

我们已经在训练数据集中对网络进行了 2 次训练。但是我们需要检查网络是否学到了什么。

我们将通过预测神经网络输出的类标签并根据实际情况来进行检查。如果预测正确，则将样本添加到正确预测列表中。

好的，第一步。让我们显示测试集中的图像，增加对其的了解。

dataiter = iter(testloader)
images, labels = dataiter.next()

# 打印图片
imshow(torchvision.utils.make_grid(images))
print("GroundTruth: ", " ".join(f"{classes[labels[j]]:5}" for j in range(4)))

接下来，让我们重新加载保存的模型（注意：这里不需要保存和重新加载模型，我们只是为了说明如何这样做）：

net = Net()
net.load_state_dict(torch.load(path))

<All keys matched successfully>

好的，现在让我们看看神经网络对以上这些示例的观点：

outputs = net(images)
outputs

tensor([[ 0.2116, -0.5072,  0.6409,  1.5528, -0.7574, -0.1588,  0.1422, -1.5925,
          1.1673, -1.4644],
        [ 4.6231,  4.3018,  0.0112, -1.9630, -2.1874, -4.9400, -4.0167, -4.6405,
          6.3403,  2.5920],
        [ 2.4601,  2.0820,  0.3965, -0.9645, -0.9624, -2.6551, -2.5377, -2.3208,
          3.2436,  0.8831],
        [ 3.9852,  0.2098,  1.7740, -1.0856,  0.4054, -3.4278, -2.6763, -2.8965,
          3.9335, -0.2969]], grad_fn=<AddmmBackward>)

输出是 10 个类别的能量。一个类别的能量越高，网络就认为该图像属于该类别。因此，让我们获取最高能量的索引：

_, predicted = torch.max(outputs, 1)
print("Predicted: ", " ".join(f"{classes[predicted[j]]:5}" for j in range(4)))

Predicted:  cat   ship  ship  plane

结果似乎相当不错。

让我们来看看网络对整个数据集的性能。

correct = 0
total = 0
with torch.no_grad():
    for data in testloader:
        images, labels = data
        outputs = net(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: 52.72%

看起来比偶然更好，准确率是 10％（从10个类别中随机选择一个类别）。好像网络学到了一些东西。

嗯，哪些类别的表现良好，哪些类别的表现不佳：

class_correct = list(0. for i in range(10))
class_total = list(0. for i in range(10))
with torch.no_grad():
    for data in testloader:
        images, labels = data
        outputs = net(images)
        _, predicted = torch.max(outputs, 1)
        c = (predicted == labels).squeeze()
        for i in range(4):
            label = labels[i]
            class_correct[label] += c[i].item()
            class_total[label] += 1
            
for i in range(10):
    print(f"Accuracy of {classes[i]:5} : {100 * class_correct[i] / class_total[i]:2.0f} %")

Accuracy of plane : 54 %
Accuracy of car   : 75 %
Accuracy of bird  : 51 %
Accuracy of cat   : 32 %
Accuracy of deer  : 36 %
Accuracy of dog   : 36 %
Accuracy of frog  : 79 %
Accuracy of horse : 44 %
Accuracy of ship  : 75 %
Accuracy of truck : 44 %

好的，那接下来呢？

我们如何在GPU上运行这些神经网络？

在 GPU 上训练

就像将 Tensor 转移到 GPU 上一样，您也将神经网络转移到 GPU 上。

如果可以使用 CUDA，首先将我们的设备定义为第一个可见的 cuda 设备：

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

# 假设我们在 CUDA 机器上，则应打印 CUDA 设备：

print(device)

cuda:0

本节的其余部分假定该 device 是 CUDA 设备。

然后，这些方法将递归遍历所有模块，并将其参数和缓冲区转换为 CUDA Tensor：

net.to(device)

Net(
  (conv1): Conv2d(3, 6, kernel_size=(5, 5), stride=(1, 1))
  (pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
  (fc1): Linear(in_features=400, out_features=120, bias=True)
  (fc2): Linear(in_features=120, out_features=84, bias=True)
  (fc3): Linear(in_features=84, out_features=10, bias=True)
)

请记住，您还必须将每一步的输入和目标也发送到 GPU：

inputs, labels = data[0].to(device), data[1].to(device)

与 CPU 相比，为什么我没有注意到 明显的 加速？因为您的网络真的很小。

练习

尝试增加网络的宽度（第一个 nn.Conv2d 的参数2和第二个 nn.Conv2d 的参数1 –它们必须是相同的数字），看看可以得到什么样的加速。

实现的目标：

全面了解 PyTorch 的 Tensor 库和神经网络。
训练一个小型神经网络对图像进行分类

在多个 GPU 上训练

如果您希望使用所有 GPU 看到更多的大规模加速，请查看可选章节：数据并行。

https://pytorch.org/tutorials/beginner/blitz/data_parallel_tutorial.html

参考

https://github.com/pytorch/tutorials

https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html