PyTorch教程-4：PyTorch中网络的训练与测试

笔者PyTorch的全部简单教程请访问：https://www.jianshu.com/nb/48831659

PyTorch教程-4：PyTorch中网络的训练与测试

基本原理

对于要训练的模型，首先我们需要定义其结构，实例化一个用于计算Loss的loss_function和一个用于更新参数的optimizer，之后的事情就比较简单了，只要准备训练数据，然后设定训练的代数（或者停止条件）就可以进行迭代的训练。最后保存模型。对于验证的模型，只要将数据传输进训练好的模型中就能得到预测的结果，当然这个过程通常是不需要计算梯度的，所以通常在 torch.no_grad() 的条件下进行。

本节以CIFAR10的一个简单训练为例进行说明

准备数据

PyTorch中提供了很好的数据接口，当然也可以配合其他的数据集与数据加载工具。对于本节中的图像类型的数据，PyTorch的 torchvision 模块提供了很好的帮助，以及 torchvision.transforms 则对图像的预处理与变换提供了很方便的方法。这些在后边的内容中会详细地介绍，这里只给出一个简单的加载数据的例子，并将CIFR10的数据分成训练集和测试集：

import torch
import torchvision
import torchvision.transforms as transforms

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

PyTorch会自动为我们下载数据集，存储在 ./data 的目录下：

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

这样我们的训练集和测试集就准备好了，后边可以直接使用它们。

定义网络结构

这里我们对上一节中的网络结构稍作更改，最主要的部分在于输入层，因为CFIAR10的数据是三通道的图片，所以输入层的通道数需要改为3，直接给出代码：

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

定理损失函数与优化器

这里也是上一节的内容，由于是多分类任务，我们使用交叉熵损失函数（Cross-Entropy Loss Function）和带动量（momentum）的SGD：

import torch.optim as optim

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

训练网路

接下来，我们开始进行网络的训练内容，其实很简单，只要写一个循环，然后按照batch从训练数据集中取出数据，然后一直喂给网络，得到网络的预测结果后，使用 loss_function 来计算损失，然后再将损失反向传播得到所有参数的梯度，最后使用 optimizer 更新参数即可。

下边的例子中，我们训练两代。trainloader（以及后边要用的testloader）都是使用PyTorch的DataLoader实例化的可迭代对象，所以可以对其进行enumerate操作：

注：python中的enumerate接收一个可迭代对象为参数（还可以有start参数用于表明起始下标，即下面的0），返回一个元组，元组第一位为迭代数（从0开始），第二位即为可迭代对象的一项：

seasons = ['Spring', 'Summer', 'Fall', 'Winter']
print(list(enumerate(seasons)))

---

[(0, 'Spring'), (1, 'Summer'), (2, 'Fall'), (3, 'Winter')]

一个例子：

epochs = 2

# running_loss to record the losses
running_loss = 0.0

for epoch in range(epochs):
    for i,data in enumerate(trainloader,0):
        # get input images and their labels
        inputs, labels = data
        # set optimizer buffer to 0
        optimizer.zero_grad()
        # forwarding
        outputs = net(inputs)
        # computing loss
        loss = loss_function(outputs, labels)
        # loss backward
        loss.backward()
        # update parameters using optimizer
        optimizer.step()

        # printing some information
        running_loss += loss.item()
        # for every 1000 mini-batches, print the loss
        if i % 1000 == 0:
            print("epoch {} - iteration {}: average loss {:.3f}".format(epoch+1, i+1, running_loss/1000))
            running_loss = 0.0


print("Training Finished!")

---

epoch 1 - iteration 0: average loss 0.001
epoch 1 - iteration 1000: average loss 1.232
epoch 1 - iteration 2000: average loss 1.258
…
epoch 1 - iteration 12000: average loss 1.156
epoch 2 - iteration 0: average loss 0.554
epoch 2 - iteration 1000: average loss 1.100
epoch 2 - iteration 2000: average loss 1.113
…
epoch 2 - iteration 12000: average loss 1.093
Training Finished!

测试网络

对训练好的网络进行测试也很简单，其实就是将待测试的数据输入给已经训练好的网络得到预测的结果即可。与训练网络不同的是，由于不需要计算梯度，所以测试网络的代码通常在 torch.no_grad() 下完成。

在本例中，由于是一个十分类任务，所以网络的输出实际上是一个大小为10的向量，表示对于该图片属于哪个类的概率的预测，所以需要对输出向量做max操作得到取得最大值的类。比如下边的代码就给出了对于测试集计算准确率的例子：

correct = 0
total = 0

with torch.no_grad():
    for data in testloader:
        images, labels = data
        outputs = net(images)
        # predicted is the class where outputs get the maximal value
        _, predicted = torch.max(outputs.data,1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print("Accuracy of the trained network over test set is {:.3f}%".format(correct/total*100))

---

Accuracy of the trained network over test set is 60.420%

在GPU上训练/测试网络

对于更大更复杂的网络，我们希望能够使用GPU进行训练或者测试，在PyTorch中，进行运算的对象必须在同一个预算设备上（CPU或者GPU，或者多块GPU的同一块上），很简单，只要将训练的网络和数据都发送到GPU上即可，这跟将一个Tensor发送到GPU上是一致的，只不过对于模型来说，只能使用to函数来传递：

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

---

cuda

这表示我们的GPU是可用的，使用to就可以将模型和数据发送到GPU上，对上述的代码进行如下的改写即可（尤其不要忘记数据的发送）：

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