已有能正常收敛的 CUDA 训练循环时，可以尝试 PyTorch 的自动混合精度：前向与损失计算放在 torch.autocast 内，用 torch.amp.GradScaler 处理 float16 梯度。它可能降低部分运算开销，但是否提速、是否节省显存以及数值是否稳定，取决于模型、GPU、算子和输入；本文没有实测性能结果。
先看适用前提
先确认当前 PyTorch 环境能使用 CUDA，且原有 float32 训练能完成一次前向、反向和参数更新。下例按 PyTorch 2.14 的 AMP 官方文档编写，旧版接口或其他设备请以对应版本文档为准。官方已提示旧的 torch.cuda.amp.autocast 等接口应改用 torch.amp 形式。
把 AMP 接进现有训练循环
假设你已经有 model、optimizer、loss_fn 和返回 inputs, targets 的 loader，将训练步改成：
import torch
device = "cuda"
model.to(device)
scaler = torch.amp.GradScaler("cuda")
model.train()
for inputs, targets in loader:
inputs = inputs.to(device)
targets = targets.to(device)
optimizer.zero_grad(set_to_none=True)
with torch.autocast(device_type="cuda", dtype=torch.float16):
outputs = model(inputs)
loss = loss_fn(outputs, targets)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
print(float(loss.detach()))
官方文档建议把前向和损失计算放在 autocast 区域，反向传播放在区域外；float16 训练常搭配 GradScaler。不要为了“混合精度”而手动把整个模型和输入都改成 half()。
怎样判断有没有帮助
先用相同数据、批大小与训练步数保存 float32 基线，再分别记录两种运行的每步耗时、峰值显存、损失曲线和最终验证指标。至少做预热后多次测量，避免把单次启动时间当性能结论。AMP 版本若出现非有限损失、指标明显恶化或算子报错，停止比较，回退基线并检查不支持的操作与精度敏感部分。
这段代码是接入模板，不是可独立训练的完整项目；没有给出具体模型、GPU 或数据集，因而不能保证加速倍数或显存节省比例。读者下一步是在自己的训练脚本中先跑一个小批次，确认损失为有限值、优化器能更新参数，再做有记录的基线对照。
Ai菜鸟网。发布者：AI小管家，转载请注明出处：https://www.alyyhw.com/29869.html
