怎样做混合精度训练

混合精度训练（Mixed Precision Training）是一种在训练深度学习模型时，同时使用 FP16（半精度浮点数） 和 FP32（单精度浮点数） 的技术，以提升训练速度、减少显存占用，同时尽可能保持模型精度。

下面从 原理 → 实现方式 → 具体代码示例 → 注意事项 给你一个完整的讲解。

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

1. 为什么用 FP16？

✅ 优点：

• 显存占用减少约 一半
• 计算速度更快（在支持 TensorCore 的 GPU 上）
• 可训练更大的 batch size 或更大的模型

❌ 缺点：

• FP16 动态范围小，容易出现：
• 梯度下溢（变成 0）
• 数值不稳定

2. 如何克服 FP16 的问题？

混合精度训练采用以下策略：

二、核心技术点

1. 损失缩放（Loss Scaling）

将 loss 放大一个倍数（如 1024、256），使得梯度在 FP16 下不至于下溢。

PyTorch 自动完成这一步。

2. 主权重（Master Weights）

模型参数在 FP32 中保留一份，用于真正的参数更新。

三、使用 PyTorch 做混合精度训练（推荐）

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

PyTorch 1.6+ 原生支持。

✅ 基本流程

from torch.cuda.amp import autocast, GradScaler
import torch
import torch.nn as nn
import torch.optim as optim
model = Model().cuda()
optimizer = optim.Adam(model.parameters())
criterion = nn.CrossEntropyLoss()
scaler = GradScaler()
for x, y in dataloader:
x, y = x.cuda(), y.cuda()
optimizer.zero_grad()
# 前向传播（自动选择 FP16 / FP32）
with autocast():
outputs = model(x)
loss = criterion(outputs, y)
# 反向传播 + 梯度缩放
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

2. autocast 的作用

with autocast():
...

• 自动决定哪些操作使用 FP16
• 哪些操作仍保持 FP32（如 softmax、BatchNorm）

3. GradScaler 的作用

scaler = GradScaler()

• 自动进行 loss scaling
• 防止梯度变成 0
• 遇到 inf / NaN 自动调整 scale

四、完整训练模板（可直接用）

for epoch in range(epochs):
model.train()
for x, y in train_loader:
x, y = x.cuda(), y.cuda()
optimizer.zero_grad()
with torch.cuda.amp.autocast():
out = model(x)
loss = criterion(out, y)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

五、使用 Hugging Face Transformers（NLP 常用）

方法一：Trainer 自动混合精度

training_args = TrainingArguments(
output_dir="./output",
fp16=True,
)

✅ 自动启用混合精度

✅ 不需要手动写 amp

六、TensorFlow / Keras 混合精度

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

⚠️ 注意：

• 最后一层输出建议用 float32
• loss 计算用 float32

七、什么时候不适合混合精度？

❌ 不适合：

• 数值极其敏感的模型
• 非常小的模型（收益不明显）
• 不支持 TensorCore 的老 GPU

✅ 非常适合：

• Transformer / LLM
• CNN（ResNet、YOLO）
• 大模型训练

八、常见坑 & 排查

九、一句话总结

混合精度训练 = FP16 计算 + FP32 参数更新 + Loss Scaling
如果你愿意，我可以：

• ✅ 帮你把现有训练代码改成混合精度
• ✅ 对比 FP32 vs FP16 的显存和速度
• ✅ 针对你用的是 CV / NLP / 大模型给你定制方案

你现在是 PyTorch 还是 TensorFlow？训练什么模型？