用 PyTorch 搭建 CNN：图像分类入门

第 20 章聊了 CNN 的原理：卷积提特征，池化降尺寸。这章就动手，把前面学的数据集、模型、损失函数、优化器、训练循环串成一条流水线。全流程走一遍，你会发现图像分类没那么神秘。

认识 CIFAR-10 数据集

CIFAR-10 是图像分类的经典入门数据集。它有 10 个类别：飞机、汽车、鸟、猫、鹿、狗、青蛙、马、船、卡车。训练集 5 万张，测试集 1 万张。

每张图是 3×32×32 的彩色图：3 个颜色通道，每个通道 32×32 像素。图很小，训练快，非常适合入门。

加载数据用 torchvision，它会自动下载数据集。我们做两步处理：先转成张量，再把像素值归一化到 [-1, 1]。

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=64, shuffle=True)

testset = torchvision.datasets.CIFAR10(
    root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(
    testset, batch_size=64, shuffle=False)

classes = ('plane', 'car', 'bird', 'cat', 'deer',
           'dog', 'frog', 'horse', 'ship', 'truck')

搭一个小型 CNN

模型是经典的两层卷积加三层全连接。卷积负责提特征，全连接负责分类。这套结构可以说是 LeNet 的现代简化版。

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


class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)   # 3 通道进，6 通道出，5×5 卷积核
        self.pool = nn.MaxPool2d(2, 2)    # 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()

注意 fc1 的输入是 16×5×5 = 400。这个数不是拍脑袋来的，是数据一层层流过去算出来的。

跟着数据流算形状

图像进网络后尺寸怎么变？自己跟着算一遍，以后搭网络就不容易懵。

1. 输入 32×32。
2. 过 5×5 卷积（无 padding），变成 28×28。
3. 过 2×2 池化，变成 14×14。
4. 再过 5×5 卷积，变成 10×10。
5. 再过 2×2 池化，变成 5×5。

此时特征图有 16 张，每张 5×5，展平就是 16×5×5 = 400 个数。这就是 fc1 输入 400 的来历。

训练：五步循环

损失函数用交叉熵 CrossEntropyLoss，它内部自带 softmax，多分类的标配。优化器用带动量的 SGD。

import torch.optim as optim

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

for epoch in range(2):  # 先跑 2 轮感受一下流程
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data

        optimizer.zero_grad()              # 1. 清空梯度
        outputs = net(inputs)              # 2. 前向传播
        loss = criterion(outputs, labels)  # 3. 算损失
        loss.backward()                    # 4. 反向传播
        optimizer.step()                   # 5. 更新参数

        running_loss += loss.item()
        if i % 200 == 199:
            print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 200:.3f}')
            running_loss = 0.0

print('Finished Training')

每 200 个小批次打印一次平均损失。看到 loss 从 2 点几一路往下掉，训练就在起作用了。

测试：准确率是多少

测试时用 torch.no_grad() 关掉梯度计算，省显存也更快。

correct = 0
total = 0
with torch.no_grad():
    for data in testloader:
        images, labels = data
        outputs = net(images)
        _, predicted = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'Accuracy: {100 * correct / total:.1f}%')

10 个类别，闭眼乱猜的准确率是 10%。上面这个小网络跑 2 轮后通常能到 50% 以上，说明它真学到了东西。想更高？加训练轮数、调大卷积通道数都是路子。

搬到 GPU 上跑

GPU 训练只需要改三处：定义 device、把模型搬过去、把每批数据搬过去。

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
net.to(device)

# 训练循环里，拿到数据后立刻搬
inputs, labels = inputs.to(device), labels.to(device)

本章的流程就是图像分类的骨架：加载数据、定义模型、选损失和优化器、训练、测试。后面学迁移学习、目标检测，都是在这个骨架上加东西。