混合精度训练怎么操作

混合精度训练（Mixed Precision Training）是指在神经网络训练过程中，同时使用 FP16（半精度）和 FP32（单精度），以加快训练速度、减少显存占用，同时尽量保持模型精度。下面从「原理 → 常见操作方式 → 框架示例 → 注意事项」系统讲一下。

一、为什么用混合精度

• ✅ 显存占用降低（FP16 是 FP32 一半）
• ✅ 计算更快（Tensor Core / GPU 半精度单元）
• ❌ 纯 FP16 容易：
• 梯度下溢（变成 0）
• 数值不稳定

所以采用混合精度：

• 前向 / 反向：FP16
• 参数主副本 / 梯度更新：FP32

二、混合精度的核心操作（通用流程）

不管什么框架，思路都类似：

1. 模型参数保留 FP32 主副本
2. 前向传播用 FP16
3. 反向传播用 FP16 算梯度
4. 梯度转回 FP32 更新参数
5. 使用 Loss Scaling 防止梯度下溢

Loss Scaling（关键）

• 小梯度 × scale（如 1024）
• 反向后 ÷ scale
• 防止 FP16 下溢为 0

三、PyTorch 中的混合精度（最常用）

1️⃣ 使用 torch.cuda.amp（推荐）

from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
optimizer.zero_grad()
with autocast():
outputs = model(inputs)
loss = criterion(outputs, labels)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

✅ 自动：

• 选择哪些层用 FP16
• Loss Scaling
• 类型转换

2️⃣ 常见参数设置

torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

四、TensorFlow / Keras

1️⃣ 开启混合精度策略

from tensorflow.keras.mixed_precision import set_global_policy
set_global_policy('mixed_float16')

2️⃣ 注意事项

• 输出层用 dtype='float32'
• optimizer 自动支持

五、Hugging Face Transformers

training_args = TrainingArguments(
fp16=True,
# bf16=True,  # A100 / 新 GPU
)

六、BF16 vs FP16（重要）

with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
...

七、常见坑

❌ 不加 Loss Scaling → 梯度为 0
❌ Embedding / Softmax 用 FP16 → 不稳定

❌ 老 GPU（无 Tensor Core）→ 提速有限

八、适用场景

✅ 大模型
✅ 显存不够
✅ 多卡训练

❌ 极小模型（收益低）

如果你愿意，我可以：

• 给你 完整训练脚本模板
• 讲 AMP 原理（数值层面）
• 针对 你用的框架 / GPU 型号 给建议

你现在是用的 PyTorch 还是 TF？GPU 是什么？