Pytorch-训练模型

上一节已经介绍了怎么定义网络，计算损失函数以及更新权重，这一节将介绍如何开始训练模型。

What about data?

当你处理图像、文本、音频、视频数据时，你可以使用标准的python包加载数据，然后将其转换成torch.Tensor。

• 对于图像，常用包有Pillow，Opencv
• 对于音频，常用包有scipy，librosa
• 对于文本，基础python，Cython，NLTK和SpaCy都很常用。
  对于视觉，有一个专门的库叫torchvision，它为常用数据集：ImageNet，CIFAR10，MNIST等提供了数据加载，还为图像提供了数据转换。

在这篇教程中，我们会使用CIFAR10数据集，它有‘airplane’, ‘automobile’, ‘bird’, ‘cat’, ‘deer’, ‘dog’, ‘frog’, ‘horse’, ‘ship’, ‘truck'类别图像，图像大小为3x32x32。

Training an image classifier

我们将会实现以下步骤：

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

1. 加载并归一化CIFAR10
   使用torchvision，加载CIFAR10将会变得非常简单：

import torch
import torchvision
import torchvision.transforms as transforms

# load CIFAR10 data and normalize
transform = transforms.Compose([transforms.toTensor(),
                                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
# train dataset
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)

# test dataset
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)

运行结果：

Downloading https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz to ./data\cifar-10-python.tar.gz
100.0%Extracting ./data\cifar-10-python.tar.gz to ./data
Files already downloaded and verified

torchvision数据集的输出是PILImage图像，像素值范围[0，1]。上面的transform做了两个操作，一个是转成Tensor，另一个是归一化像素值到[-1，1]。

如果在Windows上运行并输出BrokenPipeError，可以尝试设置torch.utils.data.DataLoader() 的参数num_worker=0。

2. 定义卷积神经网络
   模型代码如下：

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

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

4. 训练网络

# training
for epoch in range(2):
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        input, labels = data  # get (input, label)
        optimizer.zero_grad()  # set zero gradient
        output = net(input)  # forward
        loss = criterion(output, labels)  # loss
        loss.backward()  # back prop
        optimizer.step()  # update net parameters

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

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

输出结果：

[1,  2000] loss: 2.200
[1,  4000] loss: 1.866
[1,  6000] loss: 1.678
[1,  8000] loss: 1.580
[1, 10000] loss: 1.527
[1, 12000] loss: 1.461
[2,  2000] loss: 1.374
[2,  4000] loss: 1.364
[2,  6000] loss: 1.339
[2,  8000] loss: 1.324
[2, 10000] loss: 1.323
[2, 12000] loss: 1.272
Finish Training![test.png](https://img.haomeiwen.com/i11138240/859f6c600d26784a.png?imageMogr2/auto-orient/strip%7CimageView2/2/w/1240)

5. 测试网络
   我们已经用训练数据集训练了两遍网络，但是我们还需要测试网络是否学习到了分类特征。
   第一步，创建一个迭代器，并获取测试数据集的数据：

# test
dateiter=iter(testloader)
image,label=dateiter.next()

# print image
imshow(torchvision.utils.make_grid(image))
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))

# output & prediction
output = net(image)
_, predict = torch.max(output, 1)
print('Predicted: ', ' '.join('%5s' % classes[predict[j]]
                              for j in range(4)))

输出结果：

Predicted:   bird  ship  ship plane

对比可以发现，4组数据有三组正确分类。
进一步分析模型对各类型物体分类的准确性：

    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:
            image, labels = data
            outputs = net(image)
            _, 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('Accuracy of %5s : %2d %%' % (
            classes[i], 100 * class_correct[i] / class_total[i]))

输出结果：

Accuracy of plane : 74 %
Accuracy of   car : 75 %
Accuracy of  bird : 43 %
Accuracy of   cat : 27 %
Accuracy of  deer : 42 %
Accuracy of   dog : 33 %
Accuracy of  frog : 70 %
Accuracy of horse : 62 %
Accuracy of  ship : 75 %
Accuracy of truck : 43 %

在GPU上训练

首先要判断设备是否支持CUDA：

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

然后将之前的程序做一点改动：

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

这部分程序我并没有在电脑上跑，因为电脑性能的原因，我没有安装GPU版本的libtorch和pytorch，目前所有的程序都是在CPU上运行的。希望大家看完觉得写得还行的，支援我个两三毛钱，助我早日换个电脑。

结语

本节代码可以从github上下载，github上提供了cpp和python两个版本的入门程序，仅供大家参考。
C++版本程序需要的数据可以从百度云下载，提取码：vxxn