混合精度训练步骤是什么

混合精度训练（Mixed Precision Training）是一种通过使用半精度（FP16）和单精度（FP32）相结合的方式来加速深度学习模型训练的技术，同时尽量保持数值稳定性和模型精度。以下是混合精度训练的一般步骤：

一、混合精度训练的核心思想

• 前向传播：大部分计算使用 FP16（速度快、显存占用少）
• 关键步骤保留 FP32：如损失计算、参数更新等，避免数值下溢或精度损失
• 使用损失缩放（Loss Scaling）：防止 FP16 下梯度过小而变成 0

二、混合精度训练的基本步骤

1. 初始化模型和优化器

使用标准的模型结构和优化器，但后续会配合混合精度工具。

model = MyModel()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

2. 启用自动混合精度（AMP）

在 PyTorch 中通常使用 torch.cuda.amp。

from torch.cuda.amp import autocast, GradScaler

3. 创建 GradScaler（用于损失缩放）

scaler = GradScaler()

4. 训练循环中的混合精度流程

✅ 前向传播（FP16）

使用 autocast() 上下文管理器：

for data, target in dataloader:
data, target = data.cuda(), target.cuda()
optimizer.zero_grad()
with autocast():  # 自动混合精度
output = model(data)
loss = criterion(output, target)

✅ 反向传播（带损失缩放）

scaler.scale(loss).backward()

• 将 loss 放大，避免梯度下溢
• 在 FP16 下安全进行反向传播

✅ 参数更新

scaler.step(optimizer)
scaler.update()

• scaler.step() 会自动处理 unscaling
• scaler.update() 根据梯度情况动态调整缩放因子

5. 完整训练循环示例

for epoch in range(epochs):
for data, target in dataloader:
data, target = data.cuda(), target.cuda()
optimizer.zero_grad()
with autocast():
output = model(data)
loss = criterion(output, target)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

三、关键点总结

四、适用场景与注意事项

✅ 适合：

• GPU 训练（尤其是 NVIDIA Tensor Core）
• 大模型、大 batch size
• 显存受限场景

⚠️ 注意：

• 某些操作（如 BatchNorm）仍可能使用 FP32
• 不建议在 CPU 上使用混合精度
• 需要 GPU 支持 FP16（如 NVIDIA V100 / A100 / RTX 30 系列）

如果你需要 TensorFlow 的混合精度步骤 或 具体框架（如 HuggingFace、MMDetection）示例，也可以告诉我。