PyTorch 入门之官方文档学习笔记（二）训练分类器

目录

1 训练图像分类器

1.1 加载并标准化 CIFAR10 数据集

1.2 定义卷积神经网络

1.3 定义损失函数和优化器

1.4 训练网络

1.5 在测试数据上评估网络性能

2 在 GPU 上训练

---

本篇是基于官方文档做的学习笔记，在完整保留文档技术要点的同时提供辅助备注，也可作为官方文档的中文参考译文阅读。

基于上一篇，我们已经了解了如何定义神经网络、计算损失并更新网络权重。那么，数据怎么处理呢？

处理图像、文本、音频或视频数据，通常可以使用标准的 Python 包将数据加载到 NumPy 数组中，然后再将这个数组转换为 torch.*Tensor。

• 图像：可以使用 Pillow、OpenCV 等工具包
• 音频：推荐使用 scipy、librosa
• 文本：基于原生 Python 或 Cython 的加载方式，或者使用 NLTK、SpaCy

针对视觉任务，PyTorch 专门提供了一个名为 torchvision 的扩展包，其中包含常用数据集（如ImageNet、CIFAR10、MNIST等）的数据加载器，及图像数据转换工具等。

本教程将使用 CIFAR10 数据集，其包含以下类别：‘airplane’, ‘automobile’, ‘bird’, ‘cat’, ‘deer’, ‘dog’, ‘frog’, ‘horse’, ‘ship’, ‘truck’。CIFAR10 中的图像尺寸为3x32x32，即32x32像素的3通道RGB彩色图像。

1 训练图像分类器

我们将按顺序执行以下步骤：

• 使用 torchvision 加载并标准化 CIFAR10 训练集和测试集。
• 定义卷积神经网络结构。
• 定义损失函数。
• 在训练数据上训练网络。
• 在测试数据上评估网络性能。

常见的神经网络类型有：

• 前馈神经网络：数据单向流动，从输入层到输出层，无反馈连接。
• 卷积神经网络：适用于图像处理，使用卷积层提取空间特征。
• 循环神经网络：适用于序列数据，如时间序列分析和自然语言处理，允许信息反馈循环。
• 长短期记忆网络：一种特殊的 RNN，能够学习长期依赖关系。

本例是关于图像处理的，所以选择使用卷积神经网络。

1.1 加载并标准化 CIFAR10 数据集

借助 torchvision 工具包可以非常便捷地完成 CIFAR10 数据集的加载工作（该数据集包含6万张32x32像素的彩色图像，分为10个类别，其中5万张作为训练集，1万张作为测试集）。

import torch
import torchvision
import torchvision.transforms as transforms

torchvision 数据集默认输出的是 PILImage 格式的图像，原始像素值范围为 [0, 255]，首先使用transforms.ToTensor() 将其转换为 [0, 1]（即数值除以255标准化到0~1之间），然后再通过 transforms.Normalize 将它们转换为归一化范围 [-1,1] 的张量（归一化后的数据能加速模型收敛）。

注意：如果在 Windows 上运行，并且出现了一个错误的 pipeerror，请尝试将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))])

# 定义每个批次的样本数
batch_size = 4

# 训练集加载（传入前面定义的管道）
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)

# 训练数据加载器（传入前面定义的训练集，shuffle=True 打乱数据顺序（防止模型记忆顺序））
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=2)

# 测试集加载（train=False 代表加载测试集）
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)

# 测试数据加载器（shuffle=False 测试集不需要打乱顺序）
testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=2)

# 类别标签（CIFAR-10 的 10个类别名称）
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

运行后会在控制台上看到下载 CIFAR10 数据集的百分比，直到下载完毕。

可以展示一些数据集中的训练图像看一看：

import matplotlib.pyplot as plt
import numpy as np

# 展示图片方法，入参为图片张量
def imshow(img):
    img = img / 2 + 0.5 # 反归一化计算
    npimg = img.numpy() # 转换为 numpy 数组
    plt.imshow(np.transpose(npimg, (1, 2, 0))) # 维度重排（转置为 HWC）
    plt.show() # 展示图片

# 获取一些随机的训练图片
dataiter = iter(trainloader) # 将训练数据加载器转换为迭代器
images, labels = next(dataiter) # 从迭代器中获取一批数据

# 将图像网格化并展示
imshow(torchvision.utils.make_grid(images))
# 打印标签
print(' '.join(f'{classes[labels[j]]:5s}' for j in range(batch_size)))

运行结果如下：

补充说明：

• iter() 和 next() 是 Python 的迭代器协议方法，前者是将可迭代对象转换为迭代器对象，后者是从迭代器获取下一个元素。
• np.transpose(npimg, (1, 2, 0) 维度重排是为了解决 PyTorch 张量和 Matplotlib 图像显示时的维度格式不匹配问题。

1.2 定义卷积神经网络

复制前一篇“神经网络”中的网络结构，将其修改为可处理三通道图像（之前仅支持单通道图像输入）。

import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        
        # 第一个卷积层，输入通道数3（对应 RGB 三通道），输出通道数6，卷积核5x5
        self.conv1 = nn.Conv2d(3, 6, 5)
        
        # 最大池化层，池化窗口2x2，步长2
        self.pool = nn.MaxPool2d(2, 2)
        
        # 第二个卷积层，输入通道数6（与上一层输出一致），输出通道数16，卷积核5x5
        self.conv2 = nn.Conv2d(6, 16, 5)
        
        # 第一个全连接层（线性层），输入特征数16*5*5=400（假设输入图像是32x32，经过两次池化后为5x5），输出特征数120
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        
        # 第二个全连接层，输入120维，输出84维
        self.fc2 = nn.Linear(120, 84)
        
        # 第三个全连接层（输出层），输入84维，输出10维（对应10分类任务）
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        # 第一层卷积 → ReLU激活 → 最大池化，输出尺寸变化：(batch, 3, 32, 32) → (batch, 6, 14, 14)
        x = self.pool(F.relu(self.conv1(x)))
        
        # 第二层卷积 → ReLU激活 → 最大池化，输出尺寸变化：(batch, 6, 14, 14) → (batch, 16, 5, 5)
        x = self.pool(F.relu(self.conv2(x)))
        
        # 展平操作（保留batch维度），输出形状：(batch, 16*5*5) = (batch, 400)
        x = torch.flatten(x, 1)
        
        # 第一个全连接层 → ReLU激活，输出形状：(batch, 120)
        x = F.relu(self.fc1(x))
        
        # 第二个全连接层 → ReLU激活，输出形状：(batch, 84)
        x = F.relu(self.fc2(x))
        
        # 输出层（不接激活函数，通常配合 CrossEntropyLoss 使用），输出形状：(batch, 10)
        x = self.fc3(x)
        
        return x

# 实例化网络
net = Net()

1.3 定义损失函数和优化器

我们选择使用分类交叉熵损失（Classification Cross-Entropy Loss）和带动量的随机梯度下降优化器（SGD with momentum）。损失函数是计算输出与目标的偏差，优化器是告诉模型怎么改、改多少（关于损失函数和优化器，上一篇有详细讲解）。

import torch.optim as optim

# 创建交叉熵损失函数
criterion = nn.CrossEntropyLoss()

# 创建带动量的随机梯度下降优化器
# net.parameters() 告诉优化器需要调整模型中所有可训练的权重和偏置
# lr=0.001 学习率（每次参数调整的步长），类似“你学自行车时每次扭车把的幅度”
# momentum=0.9 动量（惯性系数），帮助优化器加速收敛并减少震荡
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

1.4 训练网络

现在开始进入核心阶段。我们只需循环遍历数据迭代器，将输入数据馈送到网络中进行前向传播，并执行优化（关于前向传播，上一篇有详细讲解）。

# 遍历整个数据集多次（此处设为2次）
for epoch in range(2):
    running_loss = 0.0 # 累计损失值初始化
    for i, data in enumerate(trainloader, 0):
        # 获取输入数据，data 是一个包含 [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个 mini-batch 打印一次（从0开始计数，故判断1999）
            print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}')
            running_loss = 0.0 # 重置累计损失

print('训练完成') # 训练结束提示

运行结果如下：

保存训练模型：

PATH = './cifar_net.pth' # 指定模型保存的文件路径和名称
torch.save(net.state_dict(), PATH)

点击此处，可查看有关保存 PyTorch 模型的更多详细信息。

1.5 在测试数据上评估网络性能

我们已经对训练数据集进行了2轮完整训练，但需要验证网络是否真正学到了有效特征。

我们将通过预测神经网络输出的类标签，并将其与真实值进行比较来检查这一点。如果预测正确，则将样本添加到正确预测的列表中。

第一步，展示测试集中的示例图像以便直观观察：

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

# 打印图片
imshow(torchvision.utils.make_grid(images))
print('GroundTruth: ', ' '.join(f'{classes[labels[j]]:5s}' for j in range(4)))

运行结果如下：

接下来，重新加载保存的模型（注意：这里不需要保存和重新加载模型，这样做只是为了说明如何这样做）：

net = Net()
# weights_only=True 确保加载的文件仅包含模型参数（张量），禁止加载任何可能嵌入的代码或可执行对象
net.load_state_dict(torch.load(PATH, weights_only=True))

现在看看神经网络认为上面的这些例子是什么：

outputs = net(images)

输出结果是10个类别对应的能量值。对于某个类别来说，能量值越高，网络就越认为该图像属于这个特定类别。因此，我们需要获取最高能量值对应的类别索引：

_, predicted = torch.max(outputs, 1)
print('Predicted: ', ' '.join(f'{classes[predicted[j]]:5s}' for j in range(4)))

运行结果如下：

结果看起来不错，现在来看看网络在整个数据集上的表现如何。

correct = 0 # 初始化预测正确的样本数
total = 0 # 初始化总测试样本数

# 由于当前是测试模式（非训练），不需要计算输出梯度以节省内存
with torch.no_grad():
    for data in testloader:
        images, labels = data
        
        # 将图像输入网络计算输出（前向传播）
        outputs = net(images)
        
        # 选择能量值最高的类别作为预测结果（torch.max 返回最大值及其索引）
        _, predicted = torch.max(outputs, 1)
        
        total += labels.size(0) # 累计当前批次的样本数（通常是 batch_size）
        correct += (predicted == labels).sum().item() # 累计预测正确的样本数

# 打印网络在10,000张测试图像上的准确率（取整数百分比）
print(f'网络在10000张测试图像上的准确率: {100 * correct // total} %')

代码运行结果如下：

这看起来比随机要好得多，随机精度是10%（从10个类别中随机选择一个类别）。看来网络学习到了一些东西。那么，哪些类表现得很好，哪些类表现得不好：

# 初始化每个类别的正确预测数和总预测数统计字典
correct_pred = {classname: 0 for classname in classes} # 记录每个类别预测正确的次数
total_pred = {classname: 0 for classname in classes} # 记录每个类别的总测试样本数

# 测试阶段无需计算梯度
with torch.no_grad():
    for data in testloader:
        images, labels = data # 获取测试批次数据（图像和对应标签）
        outputs = net(images) # 通过神经网络前向传播获取输出
        _, predictions = torch.max(outputs, 1) # 获取预测类别（取能量值最高的索引）
        
        # 遍历当前批次中每个样本的标签和预测结果
        for label, prediction in zip(labels, predictions):
            if label == prediction: # 如果预测正确
                correct_pred[classes[label]] += 1 # 对应类别的正确计数+1
            total_pred[classes[label]] += 1 # 对应类别的总计数+1（无论对错）


# 打印每个类别的分类准确率
for classname, correct_count in correct_pred.items():
    accuracy = 100 * float(correct_count) / total_pred[classname] # 计算准确率（转换为百分制）
    print(f'类别 [{classname:5s}] 的准确率: {accuracy:.1f} %') # 格式化输出（类别名占5字符宽度，准确率保留1位小数）

补充说明，上面字典创建代码的等效写法：

# correct_pred = {classname: 0 for classname in classes} 的等效写法
correct_pred = {}  # 创建空字典
for classname in classes:
    correct_pred[classname] = 0  # 为每个类别添加初始值0

代码运行结果如下：

接下来，让我们看看如何在 GPU 上运行这些神经网络。

2 在 GPU 上训练

就像将张量转移到 GPU 上一样，神经网络也可以转移到 GPU 上。

如果有可用的 CUDA（NVIDIA 提供的 GPU 计算框架，深度学习的主流加速工具），首先将设备定义为第一个可见的 CUDA 设备：

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

# 假设我们在 CUDA 机器上，这将打印一个 CUDA 设备：

print(device)

现在假定设备是 CUDA 设备。然后这些方法将递归遍历所有模块并将其参数和缓冲区转换为 CUDA 张量：

net.to(device)

我们需要在每一步将输入和目标发送到 GPU：

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

为什么没有明显感到与 CPU 相比的巨大加速？因为我们的网络太小了。

可以尝试增加网络的宽度（也就是第一个 nn.Conv2d 的第2个参数和第二个 nn.Conv2d 的第1个参数——二者需要相同的数字)，看看你得到什么样的加速。

【本节完&持续更新】