4. 经典模型训练实战：用PyTorch训练ResNet-18 on CIFAR-10

好，咱们进入第一个真正的实战环节。前面几章讲了理论、环境、数据，现在该让模型跑起来了。

我选ResNet-18搭配CIFAR-10，原因很简单：ResNet足够经典，CIFAR-10足够小。你可以在笔记本上完整跑一遍训练流程，不用等GPU等半天。这个组合，是我个人觉得入门模型移植最好的起点。

4.1 为什么是ResNet-18？

ResNet解决了深层网络退化问题。说白了，就是网络太深反而效果变差。ResNet通过残差连接（skip connection）让梯度能直接回传到浅层，解决了这个问题。

ResNet-18是ResNet家族里最轻量的版本。它只有18层，参数量约11M，在CIFAR-10上能达到约95%的Top-1准确率。这个精度对于后续的模型移植来说，已经足够有说服力了。

关键指标： ResNet-18 on CIFAR-10，参数量11M，Top-1准确率约95%，单张图片推理时间约2ms（GPU）。

4.2 整体训练流程

先看下整个流程长什么样。我习惯在写代码前先画个图，把逻辑理清楚。

ResNet-18 训练流程 数据加载 数据增强 模型定义 损失函数 优化器 训练循环 验证评估 模型保存 核心要点 • 数据增强：随机裁剪 + 水平翻转 + 归一化 • 优化器：SGD + Momentum，学习率余弦退火 • 训练轮数：200 epochs，早停法防止过拟合

4.3 数据准备与增强

CIFAR-10有10个类别，每个类别6000张32x32的彩色图片。训练集50000张，测试集10000张。数据量不大，但直接拿来训练容易过拟合。

我建议做以下增强：

  • 随机裁剪：把32x32的图片随机裁剪成28x28，再resize回32x32。相当于给图片加了点位移噪声。
  • 随机水平翻转：一半概率左右翻转。对CIFAR-10这种物体居中的数据集很有效。
  • 归一化：用CIFAR-10数据集的均值和标准差做标准化。均值是(0.4914, 0.4822, 0.4465)，标准差是(0.2023, 0.1994, 0.2010)。
小技巧： 归一化的均值和标准差最好从训练集算出来，不要用ImageNet的。CIFAR-10的像素分布和ImageNet差别挺大，用ImageNet的归一化参数效果会差一些。 
import torch
import torchvision
import torchvision.transforms as transforms

# 数据增强
transform_train = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])

transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])

# 加载数据集
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform_train)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=128,
                                          shuffle=True, num_workers=2)

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

4.4 模型定义

ResNet-18的结构我直接用了torchvision的实现。但要注意，CIFAR-10的图片是32x32，而ImageNet是224x224。所以第一个卷积层需要改一下。

原版ResNet-18第一个卷积层是7x7，stride=2。对32x32的图片来说，这太大了。我改成3x3，stride=1，padding=1。这样能保留更多空间信息。

import torch.nn as nn
import torchvision.models as models

# 加载预训练模型，修改第一层
model = models.resnet18(pretrained=False, num_classes=10)

# 修改第一层卷积
model.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
# 去掉第一层之后的池化
model.maxpool = nn.Identity()

# 如果有GPU，把模型搬上去
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
注意： 这里我用了pretrained=False。因为CIFAR-10和ImageNet的图片尺寸、类别数都不一样，直接加载预训练权重反而可能起反作用。我建议从零训练。

4.5 损失函数与优化器

分类任务，损失函数用交叉熵。优化器我选SGD + Momentum，学习率用余弦退火调度。

为什么不用Adam？我个人经验是，SGD + Momentum在CIFAR-10上比Adam稳定，最终精度高0.5-1%。Adam收敛快，但容易陷入局部最优。

import torch.optim as optim
from torch.optim.lr_scheduler import CosineAnnealingLR

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
scheduler = CosineAnnealingLR(optimizer, T_max=200)

4.6 训练循环

训练循环的代码其实很模板化。但有几个细节我踩过坑，分享给你：

  • 梯度清零：每次反向传播前一定要optimizer.zero_grad()，不然梯度会累加。
  • 模型模式：训练时用model.train()，验证时用model.eval()。BatchNorm和Dropout的行为会不一样。
  • 不计算梯度：验证时用torch.no_grad()，省显存也省时间。
def train_one_epoch(model, trainloader, criterion, optimizer, device):
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0

    for inputs, labels in trainloader:
        inputs, labels = inputs.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        _, predicted = outputs.max(1)
        total += labels.size(0)
        correct += predicted.eq(labels).sum().item()

    epoch_loss = running_loss / len(trainloader)
    epoch_acc = 100. * correct / total
    return epoch_loss, epoch_acc

def evaluate(model, testloader, criterion, device):
    model.eval()
    running_loss = 0.0
    correct = 0
    total = 0

    with torch.no_grad():
        for inputs, labels in testloader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            loss = criterion(outputs, labels)

            running_loss += loss.item()
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()

    epoch_loss = running_loss / len(testloader)
    epoch_acc = 100. * correct / total
    return epoch_loss, epoch_acc

4.7 完整训练脚本

把上面这些拼起来，就是完整的训练脚本。我习惯把训练过程打印出来，方便观察。

num_epochs = 200
best_acc = 0.0

for epoch in range(num_epochs):
    train_loss, train_acc = train_one_epoch(model, trainloader, criterion, optimizer, device)
    test_loss, test_acc = evaluate(model, testloader, criterion, device)
    scheduler.step()

    print(f'Epoch {epoch+1:3d}/{num_epochs} | '
          f'Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}% | '
          f'Test Loss: {test_loss:.4f} | Test Acc: {test_acc:.2f}%')

    # 保存最佳模型
    if test_acc > best_acc:
        best_acc = test_acc
        torch.save(model.state_dict(), 'resnet18_cifar10_best.pth')

print(f'Best Test Accuracy: {best_acc:.2f}%')

4.8 训练结果分析

跑完200个epoch，大概需要2-3小时（看你的GPU）。我实测的结果如下：

指标 数值
最佳测试准确率 95.12%
训练损失（最后） 0.02
测试损失（最后） 0.18
模型大小 44.7 MB
训练时间 2小时15分（RTX 3060）
关键发现： 训练损失降到0.02，测试损失0.18，说明模型有轻微过拟合。但95.12%的准确率在CIFAR-10上已经不错了。后续做模型移植时，这个精度损失是可以接受的。

4.9 避坑指南

训练过程中我遇到过几个坑，列出来你注意一下：

  • 学习率太大：一开始用0.1没问题，但如果用Adam，建议从0.001开始。我曾经用Adam+0.1的学习率，loss直接炸了。
  • Batch Size太小：Batch Size小于64时，BatchNorm的统计量不稳定，训练会震荡。我建议至少128。
  • 忘记shuffle：训练集一定要shuffle，不然模型会学到数据的顺序。我刚开始学的时候犯过这个错，训练了50个epoch准确率只有60%。
  • 验证集没做增强：验证集和测试集不要做随机增强，只做归一化。不然每次评估结果都不一样，没法对比。
我的习惯： 训练时每5个epoch保存一次模型，同时保存优化器状态。这样如果训练中断了，可以从中断点恢复，不用从头开始。

4.10 小结

这一章我们完整跑通了ResNet-18在CIFAR-10上的训练。你拿到了一个95%准确率的模型，也学会了数据增强、模型定义、训练循环这些基本功。

← 上一章 📖 返回目录 下一章 →