快捷方式

训练分类器

就是这样。您已经了解了如何定义神经网络、计算损失以及更新网络权重。

现在您可能在想,

数据怎么办?

通常,当您需要处理图像、文本、音频或视频数据时,可以使用标准的 Python 包将数据加载到 NumPy 数组中。然后,您可以将此数组转换为 torch.*Tensor

  • 对于图像,Pillow、OpenCV 等包很有用

  • 对于音频,scipy 和 librosa 等包很有用

  • 对于文本,可以使用原始 Python 或基于 Cython 的加载,或者使用 NLTK 和 SpaCy

特别是对于视觉,我们创建了一个名为 torchvision 的包,其中包含用于 ImageNet、CIFAR10、MNIST 等常见数据集的数据加载器,以及用于图像的数据转换器,即 torchvision.datasetstorch.utils.data.DataLoader

这提供了巨大的便利,避免了编写样板代码。

在本教程中,我们将使用 CIFAR10 数据集。它包含以下类别:“飞机”、“汽车”、“鸟”、“猫”、“鹿”、“狗”、“青蛙”、“马”、“船”、“卡车”。CIFAR-10 中的图像大小为 3x32x32,即大小为 32x32 像素的 3 通道彩色图像。

cifar10

cifar10

训练图像分类器

我们将按以下步骤进行操作

  1. 使用 torchvision 加载和标准化 CIFAR10 训练和测试数据集

  2. 定义卷积神经网络

  3. 定义损失函数

  4. 在训练数据上训练网络

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

1. 加载和标准化 CIFAR10

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

import torch
import torchvision
import torchvision.transforms as transforms

torchvision 数据集的输出是范围为 [0, 1] 的 PILImage 图像。我们将它们转换为范围为 [-1, 1] 的标准化张量。

注意

如果在 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))])

batch_size = 4

trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,
                                          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=batch_size,
                                         shuffle=False, num_workers=2)

classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
Downloading https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz to ./data/cifar-10-python.tar.gz

  0%|          | 0.00/170M [00:00<?, ?B/s]
  0%|          | 459k/170M [00:00<00:37, 4.52MB/s]
  5%|4         | 8.19M/170M [00:00<00:03, 47.0MB/s]
 11%|#         | 18.0M/170M [00:00<00:02, 70.3MB/s]
 17%|#6        | 28.5M/170M [00:00<00:01, 83.9MB/s]
 23%|##2       | 38.5M/170M [00:00<00:01, 89.6MB/s]
 29%|##8       | 48.9M/170M [00:00<00:01, 94.3MB/s]
 35%|###4      | 59.0M/170M [00:00<00:01, 96.3MB/s]
 41%|####      | 69.2M/170M [00:00<00:01, 98.2MB/s]
 47%|####6     | 79.5M/170M [00:00<00:00, 99.6MB/s]
 53%|#####2    | 89.5M/170M [00:01<00:00, 99.7MB/s]
 59%|#####8    | 100M/170M [00:01<00:00, 101MB/s]
 65%|######4   | 110M/170M [00:01<00:00, 100MB/s]
 70%|#######   | 120M/170M [00:01<00:00, 99.8MB/s]
 76%|#######6  | 130M/170M [00:01<00:00, 86.8MB/s]
 82%|########1 | 139M/170M [00:01<00:00, 72.3MB/s]
 86%|########6 | 147M/170M [00:01<00:00, 69.6MB/s]
 90%|######### | 154M/170M [00:01<00:00, 66.3MB/s]
 95%|#########4| 161M/170M [00:01<00:00, 65.3MB/s]
 98%|#########8| 168M/170M [00:02<00:00, 64.8MB/s]
100%|##########| 170M/170M [00:02<00:00, 79.0MB/s]
Extracting ./data/cifar-10-python.tar.gz to ./data
Files already downloaded and verified

让我们展示一些训练图像,图一乐。

import matplotlib.pyplot as plt
import numpy as np

# functions to show an image


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


# get some random training images
dataiter = iter(trainloader)
images, labels = next(dataiter)

# show images
imshow(torchvision.utils.make_grid(images))
# print labels
print(' '.join(f'{classes[labels[j]]:5s}' for j in range(batch_size)))
cifar10 tutorial
frog  plane deer  car

2. 定义卷积神经网络

复制之前神经网络部分的神经网络,并修改它以接收 3 通道图像(而不是之前定义的 1 通道图像)。

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


class Net(nn.Module):
    def __init__(self):
        super().__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 = torch.flatten(x, 1) # flatten all dimensions except batch
        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):  # loop over the dataset multiple times

    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        # get the inputs; data is a list of [inputs, labels]
        inputs, labels = data

        # zero the parameter gradients
        optimizer.zero_grad()

        # forward + backward + optimize
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        # print statistics
        running_loss += loss.item()
        if i % 2000 == 1999:    # print every 2000 mini-batches
            print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}')
            running_loss = 0.0

print('Finished Training')
[1,  2000] loss: 2.144
[1,  4000] loss: 1.835
[1,  6000] loss: 1.677
[1,  8000] loss: 1.573
[1, 10000] loss: 1.526
[1, 12000] loss: 1.447
[2,  2000] loss: 1.405
[2,  4000] loss: 1.363
[2,  6000] loss: 1.341
[2,  8000] loss: 1.340
[2, 10000] loss: 1.315
[2, 12000] loss: 1.281
Finished Training

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

PATH = './cifar_net.pth'
torch.save(net.state_dict(), PATH)

有关保存 PyTorch 模型的更多详细信息,请参阅此处

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

我们已经对训练数据集进行了 2 次训练。但我们需要检查网络是否学到了任何东西。

我们将通过预测神经网络输出的类别标签并将其与真实标签进行比较来检查这一点。如果预测正确,我们将样本添加到正确预测列表中。

好的,第一步。让我们显示测试集中的一张图像以熟悉一下。

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

# print images
imshow(torchvision.utils.make_grid(images))
print('GroundTruth: ', ' '.join(f'{classes[labels[j]]:5s}' for j in range(4)))
cifar10 tutorial
GroundTruth:  cat   ship  ship  plane

接下来,让我们重新加载我们保存的模型(注意:此处不需要保存和重新加载模型,我们只是为了说明如何操作)。

net = Net()
net.load_state_dict(torch.load(PATH, weights_only=True))
<All keys matched successfully>

好的,现在让我们看看神经网络认为上面这些示例是什么。

outputs = net(images)

输出是 10 个类别的能量。某个类别的能量越高,网络就越认为该图像是特定类别的。因此,让我们获取最高能量的索引。

_, predicted = torch.max(outputs, 1)

print('Predicted: ', ' '.join(f'{classes[predicted[j]]:5s}'
                              for j in range(4)))
Predicted:  cat   ship  truck ship

结果看起来还不错。

让我们看看网络在整个数据集上的表现。

correct = 0
total = 0
# since we're not training, we don't need to calculate the gradients for our outputs
with torch.no_grad():
    for data in testloader:
        images, labels = data
        # calculate outputs by running images through the network
        outputs = net(images)
        # the class with the highest energy is what we choose as prediction
        _, 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: 54 %

这看起来比随机猜测好多了,随机猜测的准确率是 10%(从 10 个类别中随机选择一个)。看来网络学到了一些东西。

嗯,哪些类别表现良好,哪些类别表现不佳?

# prepare to count predictions for each class
correct_pred = {classname: 0 for classname in classes}
total_pred = {classname: 0 for classname in classes}

# again no gradients needed
with torch.no_grad():
    for data in testloader:
        images, labels = data
        outputs = net(images)
        _, predictions = torch.max(outputs, 1)
        # collect the correct predictions for each class
        for label, prediction in zip(labels, predictions):
            if label == prediction:
                correct_pred[classes[label]] += 1
            total_pred[classes[label]] += 1


# print accuracy for each class
for classname, correct_count in correct_pred.items():
    accuracy = 100 * float(correct_count) / total_pred[classname]
    print(f'Accuracy for class: {classname:5s} is {accuracy:.1f} %')
Accuracy for class: plane is 37.9 %
Accuracy for class: car   is 62.2 %
Accuracy for class: bird  is 45.6 %
Accuracy for class: cat   is 29.2 %
Accuracy for class: deer  is 50.3 %
Accuracy for class: dog   is 45.9 %
Accuracy for class: frog  is 60.1 %
Accuracy for class: horse is 70.3 %
Accuracy for class: ship  is 82.9 %
Accuracy for class: truck is 63.1 %

好的,接下来该怎么办?

如何在 GPU 上运行这些神经网络?

在 GPU 上训练

就像您将张量转移到 GPU 上一样,您将神经网络转移到 GPU 上。

如果我们有 CUDA 可用,让我们首先将我们的设备定义为第一个可见的 cuda 设备。

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

# Assuming that we are on a CUDA machine, this should print a CUDA device:

print(device)
cuda:0

本节的其余部分假设device是一个 CUDA 设备。

然后这些方法将递归遍历所有模块并将它们的 parameters 和 buffers 转换为 CUDA 张量。

net.to(device)

请记住,您还必须在每一步都将输入和目标发送到 GPU。

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

为什么我没有注意到与 CPU 相比的巨大加速?因为您的网络非常小。

练习:尝试增加网络的宽度(第一个nn.Conv2d的第 2 个参数,以及第二个nn.Conv2d的第 1 个参数 - 它们需要是相同的数字),看看你能获得什么样的加速。

目标达成:

  • 从高层次理解 PyTorch 的张量库和神经网络。

  • 训练一个小型神经网络来对图像进行分类。

在多个 GPU 上训练

如果您想使用所有 GPU 看到更巨大的加速,请查看可选:数据并行

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并获得问题的解答

查看资源