📚 课程目录第2章配置PyTorch自动混合精度（AMP）

2. 配置PyTorch自动混合精度（AMP）

任务目标

在单卡GPU上，分别用纯FP32和自动混合精度（AMP）训练同一个小型CNN模型（CIFAR-10分类），对比两者的训练耗时与最终测试准确率。AMP使用torch.cuda.amp提供的GradScaler与autocast，在保持数值稳定的前提下利用FP16加速计算。

环境依赖

Python 3.8+，PyTorch 1.10+（需支持CUDA），torchvision，CUDA可用GPU
安装：pip install torch torchvision

完整实现

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

# 固定随机种子保证可复现
torch.manual_seed(0)

# 数据加载（CIFAR-10，仅用前2000张训练/500张测试加速演示）
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)
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
trainset = torch.utils.data.Subset(trainset, range(2000))
testset = torch.utils.data.Subset(testset, range(500))
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)
testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False)

# 小型CNN模型
class SmallCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 16, 3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(32 * 8 * 8, 64)
        self.fc2 = nn.Linear(64, 10)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.pool(self.relu(self.conv1(x)))
        x = self.pool(self.relu(self.conv2(x)))
        x = x.view(x.size(0), -1)
        x = self.relu(self.fc1(x))
        return self.fc2(x)

def train_eval(use_amp):
    device = torch.device('cuda')
    model = SmallCNN().to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
    scaler = torch.cuda.amp.GradScaler(enabled=use_amp)

    model.train()
    start = time.time()
    for epoch in range(3):
        for inputs, labels in trainloader:
            inputs, labels = inputs.to(device), labels.to(device)
            optimizer.zero_grad()
            with torch.cuda.amp.autocast(enabled=use_amp):
                outputs = model(inputs)
                loss = criterion(outputs, labels)
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
    train_time = time.time() - start

    # 测试
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for inputs, labels in testloader:
            inputs, labels = inputs.to(device), labels.to(device)
            with torch.cuda.amp.autocast(enabled=use_amp):
                outputs = model(inputs)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    acc = 100.0 * correct / total
    return train_time, acc

if __name__ == '__main__':
    fp32_time, fp32_acc = train_eval(use_amp=False)
    amp_time, amp_acc = train_eval(use_amp=True)
    print(f'FP32: 耗时 {fp32_time:.2f}s, 准确率 {fp32_acc:.2f}%')
    print(f'AMP:  耗时 {amp_time:.2f}s, 准确率 {amp_acc:.2f}%')
    print(f'加速比: {fp32_time / amp_time:.2f}x')

运行方式

python amp_demo.py

验证方法与判定标准

程序正常结束并输出两行对比结果，无CUDA错误或NaN损失。
AMP模式训练耗时明显低于FP32（通常加速1.2～2倍，视GPU型号而定），且两者测试准确率差距在2个百分点以内，说明AMP在提速的同时保持了精度。
若AMP准确率显著下降（>5%），检查是否遗漏scaler.update()或未在反向传播前调用scaler.scale(loss)。

关键说明

autocast上下文自动为CUDA算子选择FP16/FP32，无需手动转换输入张量。
GradScaler通过放大loss防止FP16梯度下溢，step()内部先反缩放梯度再更新权重。
本案例使用数据子集和3个epoch仅为快速演示，实际训练需使用完整数据集并增加epoch数。