如何在Python中通过混合精度训练Mixed Precision降低显存占用？-Python教程-PHP中文网

能。PyTorch的torch.cuda.amp通过autocast+GradScaler组合实现安全混合精度训练，避免model.half()导致的梯度下溢/上溢问题，显存降40%–50%，需注意BatchNorm、Softmax等敏感操作及动态缩放适配。

如何在python中通过混合精度训练mixed precision降低显存占用？

能。PyTorch 的 torch.cuda.amp 模块让混合精度训练几乎零改造接入，实测显存可降低 40%–50%，尤其对 Transformer 类大模型效果显著。

为什么直接用 `model.half()` 会崩溃？

因为 model.half() 只是简单把所有参数转成 float16，但没处理梯度下溢（如极小梯度变成 0.0）和上溢（变成 inf 或 NaN）。BatchNorm 层、Loss 计算、Softmax 输出等对数值范围敏感的操作，在纯 FP16 下极易失效。

常见错误现象包括：

训练几轮后 loss 突然变为 nan
torch.cuda.amp.GradScaler 报 found inf or nan in gradients
模型准确率卡在随机水平，完全不收敛

`autocast` + `GradScaler` 的最小可用写法

这是 PyTorch 官方推荐的 AMP 标准组合，必须成对使用，缺一不可。

关键点：

autocast 控制前向传播中哪些算子自动降为 FP16（如 nn.Linear、nn.Conv2d），但会跳过 BatchNorm、Softmax 等不安全操作
GradScaler 在反向传播前对 loss 做动态缩放（scaler.scale(loss).backward()），再在 step() 时自动取消缩放并检查梯度有效性
优化器更新仍作用于原始 FP32 参数，GradScaler 内部已确保梯度转换安全

示例代码：

Python 3.14.2

Python 3.14.2是Python编程语言在2025年12月5日发布的稳定版本，属于3.14系列的第二个维护更新。该版本包含了18项修复，重点解决了多进程、数据类及正则表达式等模块的回归问题，并修复了CVE-2025-12084等安全漏洞。此版本标志着自由线程模式（移除GIL）正式获得官方支持，是Python发展的重要里程碑。

下载

from torch.cuda.amp import autocast, GradScaler
<p>scaler = GradScaler()
for inputs, labels in dataloader:
optimizer.zero_grad()</p><pre class="brush:php;toolbar:false;">with autocast():  # 进入 FP16 自动推断上下文
    outputs = model(inputs)
    loss = criterion(outputs, labels)

scaler.scale(loss).backward()  # loss 缩放 + 反向传播
scaler.step(optimizer)         # 梯度裁剪、更新、取消缩放
scaler.update()                # 更新下一轮的缩放因子

哪些模型层或操作需要额外注意？

不是所有模块都默认兼容 AMP，以下场景容易出问题：

nn.BatchNorm2d / nn.BatchNorm1d：FP16 下 running_mean / running_var 易失稳，建议改用 nn.SyncBatchNorm 或切换为 nn.GroupNorm
自定义 loss 函数：若含手动 torch.log、torch.exp 或极小值比较（如 eps=1e-8），需显式转为 .float() 再计算
输出层带 torch.softmax：应改用 F.log_softmax + nn.NLLLoss 组合，避免 FP16 下 softmax 数值溢出
使用 torch.einsum 或自定义 CUDA kernel：需确认其内部是否支持 FP16 输入，否则强制 fallback 到 FP32