本节提供动手操作指南,介绍如何使用 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 中的两个主要组件:autocast:这是一个上下文管理器,它可以在有益且安全的情况下,将张量操作自动转换为低精度类型(在兼容 GPU 上默认为 float16)。卷积和全连接层等操作通常会大幅加速,而其他操作(如归约)可能会保留在 float32 中以保持数值稳定性。GradScaler:由于 float16 的数值范围比 float32 小得多,反向传播期间计算的梯度可能变得非常小(下溢)并被置零,从而妨碍训练。GradScaler 通过在反向传播前将损失值 向上 缩放来帮助防止这种情况。这有效地将生成的梯度缩放到 float16 的可表示范围内。在优化器更新权重之前,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 系列及更高版本)运行这两个版本,您会注意到:训练时间缩短:AMP 版本通常比标准 FP32 版本明显更快完成。加速 1.5 倍到 3 倍或更多是常见现象,具体取决于模型架构、GPU 和批次大小。峰值内存使用降低:使用 float16 张量的操作所需的内存带宽和存储量是 float32 的一半。为反向传播存储的激活值也消耗更少内存,从而可以使用更大的批次或模型。代码改动极少:集成 AMP 只需少量额外操作:初始化 GradScaler,使用 autocast 包装前向传播,以及修改 backward() 和 optimizer.step() 调用以使用 scaler。数值稳定性:得益于 GradScaler,训练过程通常保持数值稳定,收敛情况与 FP32 基准相似。由于精度变化,您可能会看到批次间的损失值略有不同,但整体训练动态通常得以保持。这里是 FP32 和 AMP 训练典型结果的示意图:{"data": [{"x": ["FP32", "AMP"], "y": [15.8, 7.5], "type": "bar", "name": "训练时间 (秒)", "marker": {"color": "#339af0"}}, {"x": ["FP32", "AMP"], "y": [2150, 1350], "type": "bar", "name": "峰值内存 (MB)", "yaxis": "y2", "marker": {"color": "#51cf66"}}], "layout": {"title": "FP32 与 AMP 性能对比 (示意图)", "yaxis": {"title": "训练时间 (秒)"}, "yaxis2": {"title": "峰值 GPU 内存 (MB)", "overlaying": "y", "side": "right"}, "barmode": "group", "legend": {"x": 0.1, "y": -0.2, "orientation": "h"}, "margin": {"l": 50, "r": 50, "t": 50, "b": 80}}}示意性对比,显示了使用 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,您可以大幅加速训练并减少内存消耗,从而可以在现有硬件上训练更大的模型或使用更大的批次。