Pytorch学习笔记（四）—— 训练一个分类器

训练一个分类器

之前的文章中已经介绍了Pytorch的基本内容，这次将会进入实战，尝试训练一个分类器并进行测试。

开始之前

Python提供了处理各种数据的package，如处理图片时可以使用Pillow、OpenCV等；处理音频时可以使用scipy、librosa等；处理文本时可以使用基于Python或Cython的原始加载或者选择NLTK、Spacy等。
但Pytorch针对视觉处理，专门创建了一个名为torchvision的package，其中包含用于常见数据集（如Imagenet，CIFAR10，MNIST等）的数据加载器，以及用于图像（即torchvision.datasets和torch.utils.data.DataLoader）的数据转换器，这为我们带来了极大的便利。
我们将使用CIFAR10数据集， 它具有以下类别：“飞机”，“汽车”，“鸟”，“猫”，“鹿”，“狗”，“青蛙”，“马”，“船”，“卡车”。 CIFAR-10中的图像尺寸为3x32x32，即尺寸为32x32像素的3通道彩色图像。

训练分类器的步骤

我们将按以下步骤完成分类器的训练：
1.使用Torchvision加载和标准化CIFAR10训练和测试数据集
2.定义卷积神经网络
3.定义损失函数
4.根据训练数据训练网络
5.在测试数据上测试网络

数据集下载

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

可以运行上述代码在官方地址下载，但速度较慢，大约需要几个小时。这里给一个百度云，下载后在项目目录新建一个data目录解压进去。

链接：CIFAR10数据集下载地址
提取码：o6uc

导入相关包以及加载数据集

import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import time
import copy
import torch.optim as optim

MINI_BATCH = 8  # 处理数据量较大的数据集时，将数据集分成mini-batch，每次加载一个mini-batch的图片
DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')  # 如果GPU可用，使用GPU

# 使用torchvision加载数据集并归一化
# Normalize(mean, std): mean:每个色彩通道的平均值，std:每个色彩通道的方差
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=False, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=MINI_BATCH, shuffle=True, num_workers=4)
# 加载测试集
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=False, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=False, num_workers=4)

注意：如果后续训练在Windows下运行时出现BrokenPipeError错误，将上述代码中DataLoader函数的num_workers参数设置为0。

定义卷积神经网络

这里将新建一个net.py文件来定义网络模型。

# net.py
import torch
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)  # 卷积层，输入3通道，输出6通道，过滤器尺寸5*5
        self.conv2 = nn.Conv2d(6, 16, 5)  # 卷积层，输入6通道，输出16通道，过滤器尺寸5*5

        self.pool = nn.MaxPool2d(2, 2)  # 池化层，使用最大池化，尺寸2*2，步长为2

        # 全连接层:16*5*5 -> 120 -> 84 -> 10
        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 = F.relu(self.conv1(x))  # 使用relu激活函数
        x = self.pool(x)
        x = F.relu(self.conv2(x))
        x = self.pool(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

该网络包含两个卷积层，每个卷积层后接一个池化层，之后是三个全连接层。最终输出的张量维度为（1*10），表示图片属于十个类别的概率。

定义训练函数，得到最优参数

回到之前的文件，继续定义模型训练函数。注意此时应导入net文件中的Net：

from net import Net

然后定义训练函数train：

def train(model, criterion, optimizer, epochs):
    since = time.time()

    best_acc = 0.0  # 记录模型测试时的最高准确率
    best_model_wts = copy.deepcopy(model.state_dict())  # 记录模型测试出的最佳参数

    for epoch in range(epochs):
        print('-' * 30)
        print('Epoch {}/{}'.format(epoch + 1, epochs))

        # 训练模型
        r