混合精度训练如何操作

混合精度训练（Mixed Precision Training）是一种在保持模型精度的同时，利用 FP16（半精度浮点数） 和 FP32（单精度浮点数） 的组合来加速训练、降低显存占用的技术。下面从原理、关键操作、常见框架实现和注意事项四个方面说明如何操作。

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

核心思路：

• 前向和反向计算用 FP16
• 参数更新用 FP32
• 通过 Loss Scaling 防止梯度下溢

二、混合精度训练的关键操作

1. 使用 FP16 进行计算

• 模型权重、激活、梯度使用 FP16
• 减少显存占用（约 40–50%）

2. Loss Scaling（损失缩放）

FP16 存在数值范围小的问题，梯度可能变为 0。
做法：

scaled_loss = loss * scale
scaled_loss.backward()
grad = grad / scale

自动混合精度（AMP）会自动完成这一步。

3. 权重以 FP32 保存（Master Weights）

• 实际更新参数时使用 FP32
• 避免长期训练精度下降

三、主流框架中的操作方式

✅ PyTorch（推荐）

方法一：使用 torch.cuda.amp（最常用）

1️⃣ 导入

from torch.cuda.amp import autocast, GradScaler

2️⃣ 初始化

scaler = GradScaler()

3️⃣ 训练循环

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

✅ Automatic Mixed Precision（AMP）

• 自动决定哪些操作用 FP16 / FP32
• 不需要手动改模型

方法二：手动 FP16（不推荐）

model.half()
data = data.half()

⚠️ 容易数值不稳定，不推荐新项目使用。

✅ TensorFlow / Keras

使用 mixed_float16

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

或在模型中：

tf.keras.mixed_precision.set_global_policy('mixed_float16')

Keras 会自动：

• 使用 FP16 计算
• 使用 FP32 保存变量
• 自动 loss scaling

✅ DeepSpeed / Megatron（大模型常用）

{
"fp16": {
"enabled": true
}
}

或：

from deepspeed import DeepSpeedConfig

用于：

• GPT / LLaMA / 多卡训练
• 节省显存 + 提速

四、适用场景与注意事项

✅ 适合使用混合精度

• GPU 训练（V100 / A100 / 3090 / 4090）
• 模型较大
• 训练速度或显存受限

❌ 注意事项

1. BN 层

• 建议 BN 保持 FP32（AMP 会自动处理）

1. Loss 过小

• 自动 loss scaling 可缓解

1. 数值不稳定

• 减小学习率
• 检查是否有 softmax + FP16 问题

1. 推理阶段

• 可转为 FP16 推理（更快）

五、效果对比（典型）

六、一句话总结

混合精度训练 = 用 FP16 跑计算，用 FP32 更新参数，靠自动 loss scaling 保证数值稳定。
如果你愿意，我可以：

• 帮你把现有训练代码改成混合精度
• 对比 FP32 / FP16 在你模型上的差异
• 讲清楚 A100 / 3090 上的实际收益