用 PyTorch 和 LeNet 完成 CIFAR-10 图像分类

这篇记录一次完整的图像分类练习：用 PyTorch 定义一个 LeNet 风格的网络，在 CIFAR-10 上训练，再加载权重预测一张图片。代码按 2022 年的教程环境整理，但核心 API 现在仍然通用。

CIFAR-10 的图片是 32 x 32 的 RGB 图像，共有 10 个类别。它比 MNIST 稍微复杂一些，正好可以用来观察彩色图像经过两次卷积和池化后，尺寸是怎样变化的。

参考资料：

目录结构

project/
├── model.py
├── train.py
├── predict.py
└── data/

第一次运行 train.py 时，CIFAR-10 会下载到 data 目录。Windows 下先把 num_workers 设为 0，确认能运行后再考虑增加数据加载进程。

定义模型：model.py

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


class LeNet(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=5)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=5)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(32 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, num_classes)

    def forward(self, x):
        # 32 -> 28 -> 14 -> 10 -> 5
        x = self.pool1(F.relu(self.conv1(x)))
        x = self.pool2(F.relu(self.conv2(x)))
        x = torch.flatten(x, start_dim=1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        return self.fc3(x)


if __name__ == "__main__":
    model = LeNet()
    sample = torch.rand(4, 3, 32, 32)
    print(model(sample).shape)

输入尺寸的变化是 32 -> 28 -> 14 -> 10 -> 5，最后得到 32 x 5 x 5 的特征图。这里的 32 是通道数，不是 batch size。

训练模型：train.py

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

from model import LeNet


CLASSES = (
    "plane", "car", "bird", "cat", "deer",
    "dog", "frog", "horse", "ship", "truck",
)


def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
    ])

    train_set = torchvision.datasets.CIFAR10(
        root="./data", train=True, download=True, transform=transform
    )
    test_set = torchvision.datasets.CIFAR10(
        root="./data", train=False, download=True, transform=transform
    )
    train_loader = torch.utils.data.DataLoader(
        train_set, batch_size=32, shuffle=True, num_workers=0
    )
    test_loader = torch.utils.data.DataLoader(
        test_set, batch_size=100, shuffle=False, num_workers=0
    )

    model = LeNet(num_classes=len(CLASSES)).to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

    for epoch in range(2):
        model.train()
        running_loss = 0.0
        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()

        print(f"epoch {epoch + 1}: loss={running_loss / len(train_loader):.4f}")

    model.eval()
    correct = total = 0
    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            predictions = model(images).argmax(dim=1)
            total += labels.size(0)
            correct += (predictions == labels).sum().item()

    print(f"accuracy: {100 * correct / total:.2f}%")
    torch.save(model.state_dict(), "lenet-cifar10.pth")


if __name__ == "__main__":
    main()

训练循环里最容易漏掉的是 optimizer.zero_grad()。如果不清空上一轮的梯度，梯度会累积，训练结果通常会变得很奇怪。CrossEntropyLoss 接收的是模型的原始 logits，所以训练时不需要在最后额外加 softmax。

单张图片预测：predict.py

import torch
import torchvision.transforms as transforms
from PIL import Image

from model import LeNet


CLASSES = (
    "plane", "car", "bird", "cat", "deer",
    "dog", "frog", "horse", "ship", "truck",
)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
transform = transforms.Compose([
    transforms.Resize((32, 32)),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])

image = transform(Image.open("horse.jpg").convert("RGB"))
image = image.unsqueeze(0).to(device)

model = LeNet(num_classes=len(CLASSES)).to(device)
model.load_state_dict(torch.load("lenet-cifar10.pth", map_location=device))
model.eval()

with torch.no_grad():
    probabilities = torch.softmax(model(image), dim=1)[0]

index = probabilities.argmax().item()
print(CLASSES[index], probabilities[index].item())