PyTorch 学习 （4） 训练分类器

PyTorch 学习 （4） 训练分类器

来自 PyTorch 教程 Deep Learning with PyTorch: A 60 Minute Blitz > Training a Classifier

关于数据

通常我们把图像、文本、声音、视频数据加载到 numpy 的数组中，然后再将其转换为 torch 的张量

• 对于图像，通常使用 Pillow , OpenCV
• 对于音频，通常使用 scipy , librosa
• 对于文本，直接使用 python 读取，或者使用 NLTK , SpaCy

针对计算机视觉，PyTorch 有一个专门的包 torchvision ， 其中包括了常见的数据集的数据加载器， 如 ImageNet , CIFAR10 , MINIST 等，以及对于图像的数据转换器，如 torchvision.datasets , torch.utils.data.Dataloader

这给我们编写代码提供了极大的便利

这里我们使用 CIFAR10 数据集，该数据集的图片被分为了十种类别，图像尺寸是 3x32x32

训练图像分类器

执行步骤如下

1. 使用 torchvision 对训练、测试数据集进行加载和标准化
2. 定义卷积神经网络 CNN
3. 定义损失函数 Loss Function
4. 通过训练数据训练网络
5. 通过测试数据测试网络

1. 加载并标准化 CIFAR10

我们可以通过 torchvision 十分容易的加载 CIFAR10

import torch
import torthvision
import torchvision.transforms as transforms

torthvison 中的数据集输出是 [0, 1] 范围的 PILImage ，我们需要将它标准化为 [-1, 1] 范围的张量

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')

torchvision.transforms 包含了很多对于图像的标准化操作，transforms.Compose 用于合并操作，它会循环执行其输入数组中的标准化操作

transforms.ToTensor 用于把 PILImage 或者 numpy.ndarray 转化为张量

transforms.Normalize 用于实现数据归一化，参数分别规定每个通道上的均值 mean 和标准差 std

我们可以通过如下代码显示训练图像

import matplotlib.pyplot as plt
import numpy as np

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

images, labels = iter(trainloader).next()
imshow(torchvision.utils.make_grid(images))

2. 定义卷积神经网络 CNN

我们对之前定义的神经网络做一些修改，修改卷积核大小，修改为接受三通道图像

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. 定义损失函数 Loss Function 和下降规则

使用分类交叉熵损失函数和带有动量的 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(5):
    running_loss = 0.0
    for i, data in enumerate(trainloader):
        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:
            print('[{}, {}] loss: {:.3f}'.format
                  (epoch + 1, i + 1, running_loss / 2000))
            running_loss = 0.0
print('Finished Training')

保存训练模型

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

5. 测试网络

通常我们需要加载已经保存的模型，然后使用该模型对数据进行分类

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

现在我们来获取结果：

outputs = net(images)

对于分类问题，得到的结果中数值最大的那个就是网络所认为的分类结果

_, predicted = torch.max(outputs, 1)
print(np.array(classes)[predicted])

类似的，我们来获得整个测试数据集的结果以及对于每一个分类的结果

correct = 0
total = 0
class_correct = np.zeros(10)
class_total = np.zeros(10)
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()
        class_correct[labels[predicted == labels]] += 1
        class_total[labels] += 1

print('Accuracy of the network on the 10000 test images: {:2f} %'.format(
    100 * correct / total))
for i in range(10):
    print('Accuracy of {} : {:2f} %'.format(
        classes[i], 100 * class_correct[i] / class_total[i]))

在测试的时候我们并不进行反馈和更新权重，仅仅是获得结果，所以使用 torch.no_grad() 可以获得更好的性能

在 GPU 上进行训练

首先获取获取 CUDA 设备 （如果有的话）

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

把网络转移到 CUDA 设备上

net.to(device)

同样的，对于输入也要如此

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

对于很小网络并不一定能看到明显的加速