怎么开启混合精度训练

开启混合精度训练（Mixed Precision Training）的方法取决于你使用的深度学习框架。下面我按常见框架给你说明，并补充一些关键注意点。

一、PyTorch 中开启混合精度

1️⃣ 使用 torch.cuda.amp（推荐，PyTorch ≥ 1.6）

import torch
from torch.cuda.amp import autocast, GradScaler
model = model.cuda()
optimizer = torch.optim.Adam(model.parameters())
scaler = GradScaler()
for data, target in dataloader:
data, target = data.cuda(), target.cuda()
optimizer.zero_grad()
# 前向使用混合精度
with autocast():
output = model(data)
loss = loss_fn(output, target)
# 反向 + 更新
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

✅ 关键点：

• autocast()：自动选择 float16 / float32
• GradScaler：防止 float16 梯度下溢

2️⃣ 使用 PyTorch Lightning（更简单）

Trainer(
precision=16,
accelerator="gpu",
devices=1
)

或：

Trainer(precision="16-mixed")

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

1️⃣ TensorFlow 2.x（推荐）

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

然后正常写模型即可。
✅ 注意：

• 最后一层（如 softmax）建议用 float32
• 优化器通常是 tf.keras.optimizers.legacy.Adam

2️⃣ Keras 示例

tf.keras.mixed_precision.set_global_policy("mixed_float16")
model = tf.keras.Sequential([
tf.keras.layers.Dense(128, activation="relu"),
tf.keras.layers.Dense(10, dtype="float32")
])

三、Hugging Face Transformers

--fp16

或：

training_args = TrainingArguments(
fp16=True,
...
)

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

1. GPU 支持

• 需要 Tensor Core（Volta 及以上：V100 / T4 / A100 / 3090 等）

1. Loss Scaling

• PyTorch：GradScaler
• TF：自动处理

1. 数值不稳定

• 少量层可强制 float32

1. Batch Norm

• 一般保持 float32 更安全

五、收益与风险

✅ 收益：

• 显存 ↓ 30–50%
• 训练速度 ↑ 1.5–3×

⚠️ 风险：

• 收敛不稳定（少数任务）
• 需要调试学习率 / scaler

如果你告诉我：

• 使用的框架（PyTorch / TF / JAX）
• 模型类型（CNN / Transformer / LLM）
• 是否多卡训练

我可以给你更具体的配置示例。