训练分类器 - 基于 PyTorch

目前为止，我们已经掌握了如何去定义神经网络、计算损失和更新网络中的权重。

关于数据

通常来讲，当你开始处理图像、文字、音频和视频数据，你可以使用 Python 的标准库加载数据进入 NumPy 的数组中。然后你可以将其转换成 torch.*Tensor 。

• 对于图片，像 Pilow、OpenCV 这样的包很有用
• 对于音频，可以使用 SciPy 和 Librosa 包
• 对于文本，可以使用原生 Python 或 Cython 加载，也可以使用 NLTK 或 SpaCy

专门针对视觉（vision），已经有创建的一个叫 torchvision 的包，此包有对于诸如 Imagenet、CIFAR10、MNIST 等等的普通数据集的加载器的 torchvision.datasets ，也有对于图像的数据转换器的 torch.utils.data.DataLoader 。

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

对于这篇博文，我们将使用 CIFAR10 数据集。它有类：“飞机”、“汽车”、“小鸟”、“猫”、“鹿”、“狗”、“青蛙”、“马”、“轮船”、“卡车”

CIFAR-10 的数据集的尺寸大小为 \(3\times 32\times 32\) ，其中，3 个颜色通道和 \(32\times 32\) 的像素。

训练一个图像分类器

我们将按照下面的步骤：

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] 标准化的 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')

Files already downloaded and verified
Files already downloaded and verified

让我们看一看训练集中的一些图片。

import matplotlib.pyplot as plt
import numpy as np

# 显示图像的函数
def imshow(img):
    img = img / 2 + 0.5  # unnormalize
    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)))

 deer   car plane  deer

2. 定义卷积神经网络

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. 定义损失函数和优化器

让我们使用分类交叉熵损失函数和带动量的随机梯度下降法。

import torch.optim as optim

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

4. 训练网络

这时事情变得有趣起来。我们只需要遍历我们的数据迭代器，并将输入馈送到网络中并优化即可。

for epoch in range(2):  # 在数据集上多循环几次
    
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        # 得到输入和标签，数据是一个列表 [inputs, labels]
        inputs, labels = data
        
        # 将参数梯度清零
        optimizer.zero_grad()
        
        # 前向传播 + 反向传播 + 优化
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        # 输出统计数据
        running_loss += loss.item()
        if i % 2000 == 1999:  # 每 2000 为一个小批量输出
            print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 2000))
            running_loss = 0.0

[1,  2000] loss: 2.107
[1,  4000] loss: 1.965
[1,  6000] loss: 1.934
[1,  8000] loss: 1.951
[1, 10000] loss: 1.927
[1, 12000] loss: 1.963
[2,  2000] loss: 1.941
[2,  4000] loss: 1.950
[2,  6000] loss: 1.973
[2,  8000] loss: 1.980
[2, 10000] loss: 1.981
[2, 12000] loss: 1.998

现在，我们将训练完的模型存储起来。

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

可以参考官方文档来查看更多的存储 PyTorch 模型的细节。

5. 在测试集上测试

我们已经在训练集上训练了 2 遍的网络。但是，我们仍需要检查网络是否已经学得了什么。

我们将通过预测神经网络输出的类标签，并检查它是否和真实标签相同。如果预测错误，我们将出错的样本添加到错误预测的列表中。

首先，让我们输出一个测试集的图像熟练下。

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

# 输出图像
imshow(torchvision.utils.make_grid(images))
print('GroundTruth: ', ' '.join('%5s' % classes[labels[j]] for j in range(4)))

GroundTruth:    cat  ship  ship plane

下一步，让我们重新加载我们存储了的模型。

注意：保存和重新加载模型在这里不是必须的，这里只是演示一下如何去做。

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

<All keys matched successfully>

现在，让我们来看一下神经网络认为上面的图像属于哪一类

outputs = net(images)

输出是关于 10 个类的激励值。对于每个类，越高的激励值，网络就越认为这个图像属于哪一个类。所以，让我们先得到最高的激励值的索引：

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

print('Predicted: ', ' '.join('%5s' % classes[predicted[j]] for j in range(4)))

Predicted:   ship plane plane 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('神经网络在 10000 个测试图像上的准确率（Accuracy）: %d %%' % (100 * correct / total))

神经网络在 10000 个测试图像上的准确率（Accuracy）: 24 %

这看上去比 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('%5s 的准确率: %2d %%' % (classes[i], 100 * class_correct[i] / class_total[i]))

plane 的准确率: 45 %
  car 的准确率: 49 %
 bird 的准确率:  0 %
  cat 的准确率:  6 %
 deer 的准确率: 73 %
  dog 的准确率:  0 %
 frog 的准确率:  7 %
horse 的准确率: 38 %
 ship 的准确率: 13 %
truck 的准确率:  7 %

在 GPU 上训练

就像将一个张量转换进 GPU 一样，也可以将网络转进 GPU

如果我们的 CUDA 可用，首先定义我们的一块 cuda 设备：

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

# 假设我们在 CUDA 的设备上，下面将会看到输出一个 CUDA 设备

print(device)

cpu

下面剩下的部分假设 device 是一块 CUDA 设备。

然后，这些方法将递归的遍历所有模块并将它们的参数和缓冲区转换到 CUDA 张量：

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)

在多块 GPU 上训练

如果需要使用电脑上所有的 GPU 来获得更多的加速，就参考官方文档的数据并行