动手实战：实现混合精度训练

提供一个动手操作指南，介绍如何使用 torch.cuda.amp 实现自动混合精度 (AMP) 训练。自动混合精度 (AMP) 允许在特定操作中使用低精度浮点格式（如 float16）。这种能力可以大幅提升速度并减少 GPU 内存占用，通常对模型精度影响极小。该指南将讲解如何将标准的 PyTorch 训练循环转换为使用 AMP，并说明必要的变更。

本实践练习假设您可以使用支持 CUDA 的 NVIDIA GPU，其计算能力为 7.0 或更高（这是高效进行 float16 张量核心操作的必要条件），并且 PyTorch 安装版本相对较新（1.6 或更高）。

基准：标准训练循环

我们从一个使用全精度（float32）的简化标准训练循环开始。我们将使用一个基本的卷积网络和随机数据进行演示。

import torch
import torch.nn as nn
import torch.optim as optim
import time
import contextlib # 用于计时上下文管理器

# 1. 定义一个简单模型
class SimpleCNN(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
        self.relu = nn.ReLU()
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        # 为全连接层展平特征
        self.fc = nn.Linear(64 * 16 * 16, num_classes) # 假设输入图像为 32x32

    def forward(self, x):
        x = self.pool(self.relu(self.conv1(x)))
        x = self.pool(self.relu(self.conv2(x)))
        x = torch.flatten(x, 1) # 展平除批次维度外的所有维度
        x = self.fc(x)
        return x

# 2. 设置：设备、模型、数据、损失函数、优化器
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")

model = SimpleCNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)

# 模拟数据参数
batch_size = 64
img_size = 32
num_batches = 100

# 简单计时器上下文管理器
@contextlib.contextmanager
def measure_time():
    start_time = time.time()
    yield
    end_time = time.time()
    print(f"耗时: {end_time - start_time:.4f} 秒")

# 简单内存使用报告器
def report_memory(stage=""):
    if torch.cuda.is_available():
        print(f"{stage} - 峰值内存分配: {torch.cuda.max_memory_allocated(device) / 1e6:.2f} MB")
        torch.cuda.reset_peak_memory_stats(device) # 为下次测量重置峰值计数器

# 3. 标准训练循环 (FP32)
print("\n--- 标准 FP32 训练 ---")
report_memory("训练前")
model.train()
with measure_time():
    for i in range(num_batches):
        # 实时生成模拟数据
        inputs = torch.randn(batch_size, 3, img_size, img_size, device=device)
        labels = torch.randint(0, 10, (batch_size,), device=device)

        optimizer.zero_grad()

        # 前向传播
        outputs = model(inputs)
        loss = criterion(outputs, labels)

        # 反向传播和优化
        loss.backward()
        optimizer.step()

        if (i + 1) % 20 == 0:
             print(f"批次 [{i+1}/{num_batches}]，损失: {loss.item():.4f}")

report_memory("训练后")
print("--- 标准 FP32 训练完成 ---")

运行此代码（如果您有合适的 GPU）。请记录报告的耗时和峰值内存使用。这作为我们的基准。

使用 torch.cuda.amp 实现混合精度

现在，我们来修改循环以集成 AMP。这需要 torch.cuda.amp 中的两个主要组件：

1. autocast：这是一个上下文 (context)管理器，它可以在有益且安全的情况下，将张量操作自动转换为低精度类型（在兼容 GPU 上默认为 float16）。卷积和全连接层等操作通常会大幅加速，而其他操作（如归约）可能会保留在 float32 中以保持数值稳定性。
2. GradScaler：由于 float16 的数值范围比 float32 小得多，反向传播 (backpropagation)期间计算的梯度可能变得非常小（下溢）并被置零，从而妨碍训练。GradScaler 通过在反向传播前将损失值 向上 缩放来帮助防止这种情况。这有效地将生成的梯度缩放到 float16 的可表示范围内。在优化器更新权重 (weight)之前，GradScaler 会将梯度 反向缩放 回其原始值。如果在反向缩放期间检测到任何非有限梯度（NaN 或 Inf）（这有时会发生在训练不稳定或损失缩放因子较高时），则会跳过该批次的优化器步骤。GradScaler 还会随时间动态调整缩放因子。

以下是我们修改训练循环的方式：

import torch
import torch.nn as nn
import torch.optim as optim
import time
import contextlib # 用于计时上下文管理器
from torch.cuda.amp import GradScaler, autocast

# --- 重新初始化模型和优化器以进行公平比较 ---
model = SimpleCNN().to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
# --- 保持损失函数和模拟数据参数不变 ---
criterion = nn.CrossEntropyLoss()
batch_size = 64
img_size = 32
num_batches = 100

# --- 使用相同的计时器和内存报告器 ---
@contextlib.contextmanager
def measure_time():
    start_time = time.time()
    yield
    end_time = time.time()
    print(f"耗时: {end_time - start_time:.4f} 秒")

def report_memory(stage=""):
    if torch.cuda.is_available():
        print(f"{stage} - 峰值内存分配: {torch.cuda.max_memory_allocated(device) / 1e6:.2f} MB")
        torch.cuda.reset_peak_memory_stats(device)

print("\n--- 混合精度 (AMP) 训练 ---")
# 1. 初始化 GradScaler
scaler = GradScaler()

report_memory("训练前")
model.train()
with measure_time():
    for i in range(num_batches):
        inputs = torch.randn(batch_size, 3, img_size, img_size, device=device)
        labels = torch.randint(0, 10, (batch_size,), device=device)

        optimizer.zero_grad()

        # 2. 使用 autocast 包装前向传播
        # 此上下文中的操作在支持的情况下以低精度 (FP16) 运行
        with autocast():
            outputs = model(inputs)
            loss = criterion(outputs, labels)

        # 3. 在 backward() 之前缩放损失
        # scaler.scale 将损失乘以当前缩放因子
        scaler.scale(loss).backward()

        # 4. scaler.step() 反向缩放梯度并调用 optimizer.step()
        # 反向缩放位于 optimizer.param_groups[...].grad 中的梯度
        # 如果梯度不是有限的 (NaN/Inf)，则跳过 optimizer.step()
        scaler.step(optimizer)

        # 5. 为下一次迭代更新缩放因子
        # 如果发现 NaNs/Infs，则减小缩放因子，否则可能增大
        scaler.update()

        if (i + 1) % 20 == 0:
            # 注意：loss.item() 仍是未缩放的损失值
            print(f"批次 [{i+1}/{num_batches}]，损失: {loss.item():.4f}")

report_memory("训练后")
print("--- 混合精度 (AMP) 训练完成 ---")

分析与观察

如果您在兼容的 GPU 上（特别是带有张量核心的 GPU，如 V100、T4、A100、H100 或 RTX 20xx 系列及更高版本）运行这两个版本，您会注意到：

1. 训练时间缩短：AMP 版本通常比标准 FP32 版本明显更快完成。加速 1.5 倍到 3 倍或更多是常见现象，具体取决于模型架构、GPU 和批次大小。
2. 峰值内存使用降低：使用 float16 张量的操作所需的内存带宽和存储量是 float32 的一半。为反向传播 (backpropagation)存储的激活值也消耗更少内存，从而可以使用更大的批次或模型。
3. 代码改动极少：集成 AMP 只需少量额外操作：初始化 GradScaler，使用 autocast 包装前向传播，以及修改 backward() 和 optimizer.step() 调用以使用 scaler。
4. 数值稳定性：得益于 GradScaler，训练过程通常保持数值稳定，收敛情况与 FP32 基准相似。由于精度变化，您可能会看到批次间的损失值略有不同，但整体训练动态通常得以保持。

这里是 FP32 和 AMP 训练典型结果的示意图：

示意性对比，显示了使用 AMP 与标准 FP32 训练相比，典型的加速和内存减少情况。实际结果会因硬件和模型而异。

更多考量

• bfloat16：在较新的硬件（例如，A100 等 Ampere 架构 GPU，或较新的 TPU）上，您可能更喜欢使用 torch.bfloat16。它具有与 float32 相同的指数范围，但尾数精度较低。这通常使其对下溢/上溢问题更具抵御能力，并有可能无需 GradScaler。您可以通过 autocast(dtype=torch.bfloat16) 启用它。请查阅您的硬件文档以获取最佳设置。
• 性能分析：尽管 AMP 通常提供开箱即用的优势，但建议使用 PyTorch 性能分析器（在第 4 章中介绍）来确认预期加速得以实现，并找出特定于您的模型或操作、可能无法从低精度中获益的任何潜在瓶颈。
• 梯度裁剪：AMP 有时会与梯度裁剪配合使用。通常建议在裁剪梯度 之前 对其进行反向缩放。scaler.unscale_(optimizer) 可以在调用 torch.nn.utils.clip_grad_norm_ 或 torch.nn.utils.clip_grad_value_ 之前调用。

本实践练习展现了将自动混合精度集成到您的 PyTorch 训练循环中是多么容易。借助 torch.cuda.amp，您可以大幅加速训练并减少内存消耗，从而可以在现有硬件上训练更大的模型或使用更大的批次。