混合精度训练怎样实现

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

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

1. 为什么要用 FP16

• 优点
• 显存占用减半（FP16 占 2 字节，FP32 占 4 字节）
• 在支持 Tensor Core 的 GPU（如 NVIDIA V100 / A100 / 4090）上计算更快
• 问题
• 数值范围小，容易：
• 下溢（梯度变成 0）
• 上溢（梯度变成 inf / nan）

2. 混合精度的解决方案

混合精度训练通过以下机制解决 FP16 的问题：

✅ 1）权重主副本（Master Weights）

• 优化器更新使用 FP32
• 模型保存一份 FP32 的主权重
• 前向 / 反向计算使用 FP16
• 更新时：

FP16 权重 → 计算梯度 → 转成 FP32 → 更新 FP32 权重

✅ 2）Loss Scaling（损失缩放）

• 在反向传播前，将 loss 乘以一个较大的 scale（如 1024、2048）
• 防止梯度下溢
• 反向传播后：
• 梯度除以 scale
• 再更新权重

二、混合精度训练的基本流程

FP32 权重
↓
复制到 FP16
↓
前向传播 (FP16)
↓
计算 Loss
↓
Loss × scale
↓
反向传播 (FP16)
↓
梯度 / scale
↓
更新 FP32 权重

三、不同框架中的实现方式

四、PyTorch 实现（最常用）

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

示例（最简版）

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():          # 自动混合精度
y_pred = model(x)
loss = criterion(y_pred, y)
scaler.scale(loss).backward()     # loss scaling
scaler.step(optimizer)
scaler.update()

✅ 关键点

• autocast()：自动决定哪些操作用 FP16
• GradScaler：自动管理 loss scaling

2️⃣ 自定义训练循环（更灵活）

with autocast():
outputs = model(inputs)
loss = loss_fn(outputs, targets)
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(optimizer)
scaler.update()

五、TensorFlow / Keras 实现

1️⃣ 使用 mixed_float16

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

然后正常写模型即可：

model = tf.keras.Sequential([
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dense(10)
])
model.compile(
optimizer='adam',
loss='sparse_categorical_crossentropy'
)

✅ TensorFlow 会自动：

• 使用 FP16 计算
• 维护 FP32 master weights
• 处理 loss scaling

六、什么时候该用混合精度

✅ 适合

• 使用 NVIDIA GPU（Volta 架构及以上）
• 大模型 / 大 batch
• Transformer、CNN、Diffusion 等

❌ 不适合

• 数值极其敏感的模型
• 老 GPU（不支持 Tensor Core）

七、常见坑 & 注意事项

八、一句话总结

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

• 帮你把 现有训练代码改成混合精度
• 对比 FP32 vs FP16 的显存和速度
• 讲 AMP 在 LLM / Vision / Diffusion 中的实战经验

你现在是用的 PyTorch 还是 TensorFlow？