汤不热吧自动混合精度（Automatic Mixed Precision, AMP）训练是 PyTorch 1.6+ 版本引入的一项重要功能，它允许模型在训练过程中自动使用 FP16（半精度浮点数）进行计算，同时保留 FP32（单精度浮点数）来处理关键操作（如权重更新和梯度累加），从而在不损失模型精度的情况下，大幅降低显存占用并提升训练速度（通常带来 1.5x 到 3x 的吞吐量提升）。
本文将聚焦于如何简单、高效地在现有 PyTorch 训练代码中集成 torch.cuda.amp。
1. 为什么 AMP 能加速训练？
现代英伟达 GPU（如 V100, A100, RTX 系列）配备了 Tensor Core，这些核心专门用于高效执行 FP16 矩阵乘法。通过混合精度，我们能够利用 Tensor Core 的强大计算能力，同时因为 FP16 占用内存只有 FP32 的一半，我们还能在有限的显存中放入更大的 Batch Size，进一步提高 GPU 利用率。
2. 核心组件：autocast 和 GradScaler
要在 PyTorch 中实现 AMP 训练，只需要关注两个关键对象：
2.1. torch.cuda.amp.autocast
这是一个上下文管理器（Context Manager）。在 with autocast(): 内部，PyTorch 会自动选择精度最高的 dtype（通常是 FP16）来执行模型的前向传播和损失计算。如果某个操作不支持 FP16，它会自动回退到 FP32。
2.2. torch.cuda.amp.GradScaler
FP16 的数值范围比 FP32 小很多。在反向传播过程中，模型的梯度往往非常小，容易发生下溢（Underflow），导致梯度变为 0。GradScaler 通过动态梯度放大（Gradient Scaling）来解决这个问题：
1. 在反向传播前，将损失乘以一个较大的缩放因子（Scale）。
2. 计算得到的梯度也会被等比例放大，避免下溢。
3. 在执行优化器更新权重之前，GradScaler 会将梯度除以该缩放因子，恢复其原始大小，然后检查是否有 Inf/NaN 值（溢出）。只有梯度有效时，优化器才会执行 step()。
3. 实战代码示例：集成 AMP
以下代码展示了如何将标准的 PyTorch 训练流程修改为支持自动混合精度，并对比其速度。
import torch
import torch.nn as nn
import torch.optim as optim
from torch.cuda.amp import autocast, GradScaler
import time
from torchvision.models import resnet18
# --- 1. 环境准备与模型配置 ---
if not torch.cuda.is_available():
print("CUDA not available. Exiting.")
exit()
device = torch.device("cuda")
# 假设数据和模型配置
BATCH_SIZE = 128 # 尝试使用较大的 Batch Size
NUM_STEPS = 50
INPUT_SIZE = (BATCH_SIZE, 3, 224, 224)
TARGET_SIZE = (BATCH_SIZE,)
# 创建模型、优化器和数据
def setup_model():
model = resnet18(num_classes=10).to(device)
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()
input_data = torch.randn(INPUT_SIZE, device=device)
target_data = torch.randint(0, 10, TARGET_SIZE, device=device)
return model, optimizer, criterion, input_data, target_data
# --- 2. FP32 基准训练（Standard PyTorch） ---
def train_fp32(model, optimizer, criterion, input_data, target_data):
model.train()
start_time = time.time()
for _ in range(NUM_STEPS):
optimizer.zero_grad()
# 前向传播 (FP32)
output = model(input_data)
loss = criterion(output, target_data)
# 反向传播与优化 (FP32)
loss.backward()
optimizer.step()
return time.time() - start_time
# --- 3. AMP 混合精度训练 ---
def train_amp(model, optimizer, criterion, input_data, target_data):
model.train()
# 实例化 GradScaler
scaler = GradScaler()
start_time = time.time()
for _ in range(NUM_STEPS):
optimizer.zero_grad()
# 关键步骤 1: 使用 autocast 上下文管理器
with autocast():
output = model(input_data)
loss = criterion(output, target_data)
# 关键步骤 2: 替代 loss.backward()，使用 scaler.scale()
scaler.scale(loss).backward()
# 关键步骤 3: 替代 optimizer.step()，使用 scaler.step()
scaler.step(optimizer)
# 关键步骤 4: 更新缩放因子
scaler.update()
return time.time() - start_time
# --- 4. 运行与结果对比 ---
# 运行 FP32 基准
model_fp32, optimizer_fp32, criterion, input_data, target_data = setup_model()
time_fp32 = train_fp32(model_fp32, optimizer_fp32, criterion, input_data, target_data)
print(f"\nFP32 总耗时 ({NUM_STEPS}步): {time_fp32:.4f} 秒")
# 运行 AMP
model_amp, optimizer_amp, criterion, input_data, target_data = setup_model()
time_amp = train_amp(model_amp, optimizer_amp, criterion, input_data, target_data)
print(f"AMP 总耗时 ({NUM_STEPS}步): {time_amp:.4f} 秒")
if time_fp32 > 0:
speedup = time_fp32 / time_amp
print(f"吞吐量提升倍数: {speedup:.2f}X")
4. 总结与注意事项
集成 torch.cuda.amp 只需添加 GradScaler 的初始化以及在训练循环中包裹 autocast，并替换标准的梯度计算和优化器步骤。
注意事项：
1. 设备要求： 必须使用支持 FP16/Tensor Core 的 NVIDIA GPU（如 Volta、Turing、Ampere 及后续架构）。
2. DataLoader/Dataloader： 数据加载部分无需修改，输入数据仍以 FP32 格式传入 GPU，autocast 会在模型内部将其转换为 FP16。
3. 权重保存： GradScaler 的状态也是需要保存的重要信息，以确保从中断点恢复训练时，梯度缩放因子能够继续工作。使用 torch.save 时应同时保存 scaler.state_dict()。