如何开启混合精度训练

开启混合精度训练（Mixed Precision Training）的方法取决于你使用的深度学习框架。下面以 PyTorch 和 TensorFlow 为例，给出常见、实用的开启方式。

一、PyTorch 中开启混合精度训练

✅ 推荐方式：使用 torch.cuda.amp（自动混合精度）

1️⃣ 基本用法（最常用）

import torch
from torch.cuda.amp import autocast, GradScaler
model = MyModel().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 = loss_fn(y_pred, y)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

✅ 效果

• 自动在 float16 / float32 之间切换
• 减少显存占用
• 提升训练速度（尤其在 A100 / V100 / 3090 上）

2️⃣ 使用 torch.nn.parallel.DistributedDataParallel

with autocast():
output = model(input)
loss = loss_fn(output, target)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

⚠️ autocast 只需要在 forward 和 loss 计算阶段使用

3️⃣ PyTorch 2.x（可选）

PyTorch 2.x 仍推荐使用 torch.cuda.amp，尚未完全统一为 torch.compile + AMP。

二、TensorFlow / Keras 中开启混合精度

✅ 推荐方式：设置全局策略（最简单）

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, dtype='float32')  # 输出层建议 float32
])
model.compile(optimizer='adam', loss='categorical_crossentropy')

✅ 原因

• 输出层用 float32 可提高数值稳定性
• 中间层自动使用 float16

✅ 对自定义训练循环（更精细控制）

from tensorflow.keras.mixed_precision import LossScaleOptimizer
optimizer = LossScaleOptimizer(tf.keras.optimizers.Adam())

三、什么时候该用混合精度？

✅ 适合使用的情况

• 使用 GPU（尤其是 Tensor Core）
• 模型较大、显存吃紧
• 训练时间开销大

❌ 不适合的情况

• CPU 训练
• 极度数值敏感任务（如某些科学计算）

四、注意事项（非常重要）

1. Loss Scaling

• PyTorch：GradScaler
• TensorFlow：自动完成

1. BatchNorm

• 通常仍然是 float32 计算，框架自动处理

1. 评估阶段

• 推理时可以关闭 AMP 或继续开启

五、快速总结

如果你愿意，可以告诉我：

• 使用的是 PyTorch 还是 TensorFlow
• 是否是 单机 / 多机 / 多卡
• 使用的 显卡型号

我可以给你一份完全可直接复制的训练模板。