混合精度训练怎么实现

混合精度训练（Mixed Precision Training）是一种在训练深度学习模型时，同时使用 FP16（半精度）和 FP32（单精度） 的技术，目的是在几乎不损失精度的前提下，加快训练速度、减少显存占用。下面从原理、实现方式、代码示例和注意事项四个方面说明。

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

1. 为什么要用混合精度？

• FP16 优点
• 显存占用减半
• 计算更快（Tensor Core）
• FP16 缺点
• 数值范围小，容易溢出（梯度下溢 / 上溢）
• 训练不稳定

2. 核心思路

用 FP16 做大部分计算，用 FP32 保存关键数值
主要包括：

• ✅ 前向传播：FP16
• ✅ 反向传播：FP16
• ✅ 权重更新：FP32（通过 Master Weights）
• ✅ Loss 缩放（Loss Scaling）

二、关键技术点

1. 主权重（Master Weights）

• 模型参数在 FP32 中保留一份副本
• 每次更新权重时：

1. FP16 梯度 → 转为 FP32
2. 更新 FP32 权重
3. 再转回 FP16 用于下一轮前向

2. Loss Scaling

• 在反向传播前 放大 loss
• 防止梯度下溢（变成 0）
• 更新权重前再缩小回来

三、常见实现方式

1️⃣ PyTorch（推荐）

✅ PyTorch >= 1.6：torch.cuda.amp

AMP = Automatic Mixed Precision

最小可用示例

import torch
from torch.cuda.amp import autocast, GradScaler
model = model.cuda()
optimizer = torch.optim.Adam(model.parameters())
scaler = GradScaler()
for x, y in dataloader:
x, y = x.cuda(), y.cuda()
optimizer.zero_grad()
with autocast():          # FP16 自动混合精度
y_hat = model(x)
loss = criterion(y_hat, y)
scaler.scale(loss).backward()     # loss scaling
scaler.step(optimizer)
scaler.update()

✅ 优点

• 自动管理 FP16 / FP32
• 自动 loss scaling
• 几乎不用改模型代码

2️⃣ TensorFlow / Keras

✅ TF 2.x：tf.keras.mixed_precision

from tensorflow.keras import mixed_precision
policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_global_policy(policy)
model = build_model()
model.compile(optimizer='adam', loss='categorical_crossentropy')
model.fit(x, y)

✅ TF 会自动：

• 使用 float16 计算
• 用 float32 做变量更新

3️⃣ 手动实现（不推荐）

• 自行管理 master weights
• 自行实现 loss scaling
• 容错成本高

四、适用范围与注意事项

✅ 适合混合精度

• Transformer / CNN / 大模型
• NVIDIA GPU（支持 Tensor Core）
• 批量较大、训练慢的场景

⚠️ 注意事项

五、什么时候用混合精度？

✅ 推荐：

• 训练大模型
• 显存不足
• 想加速训练

❌ 不推荐：

• 非常小的模型
• 对数值精度极端敏感的任务
• 不支持 Tensor Core 的设备

六、一句话总结

混合精度训练 = FP16 计算 + FP32 主权重 + Loss Scaling
如果你愿意，我可以：

• 帮你把现有训练代码改成混合精度
• 对比 FP32 vs FP16 的显存/速度
• 结合你使用的框架（PyTorch / TF / 分布式训练）详细讲