自动混合精度 AMP
AMP 为适合操作选择低精度，并用 GradScaler 缓解 float16 梯度下溢。
学习目标
- CPU 上验证 AMP 关闭后的完整训练步
- AMP 时可以先裁剪再 unscale 吗？
- 说明：本页 CPU 代码验证回退流程，未测量或宣称 CUDA 性能收益。
为什么需要这个概念
自动混合精度由 autocast 按设备和算子选择 dtype；float16 训练常配合 GradScaler 对损失缩放，降低小梯度下溢风险。bfloat16 的指数范围不同，通常不需要同样的缩放策略。
学习前需要掌握
背景与问题
浮点精度会影响显存、吞吐和数值范围。AMP 根据算子选择适合的精度，而不是简单把模型和所有输入一律转成半精度。
概念定义
自动混合精度由 autocast 按设备和算子选择 dtype；float16 训练常配合 GradScaler 对损失缩放，降低小梯度下溢风险。bfloat16 的指数范围不同，通常不需要同样的缩放策略。
直观理解
不同工序使用不同精度的量尺：大矩阵乘法可用更便宜的尺，敏感累加与损失计算保留更稳的精度。
核心原理
autocast 通常包住前向和损失，不包住整个优化器逻辑。模型参数往往仍保留 float32 主副本；实际算子是否降精度由支持策略决定，不能从上下文名称推断每个输出 dtype。
GradScaler 先将损失乘缩放因子，再反向；在更新前还原梯度，检测非有限值并按需跳过更新、调整缩放。若需要裁剪，必须先 unscale，再按原梯度尺度裁剪。
CUDA float16、CPU bfloat16 的支持和加速效果不同。CPU 示例可以明确关闭 AMP，验证同一训练路径的回退；它不能证明 GPU 加速，也不能在没有 CUDA 的机器上伪造吞吐数字。
数学表达
核心公式
自动混合精度由 autocast 按设备和算子选择 dtype；float16 训练常配合 GradScaler 对损失缩放，降低小梯度下溢风险。bfloat16 的指数范围不同，通常不需要同样的缩放策略。
变量含义
Nbatch sizeD/H特征或隐藏维度
计算步骤
- 1前向平方损失为 16。
- 2scaler 在关闭状态下透传操作。
- 3更新后权重为 2.6。
代码对应位置
示例 1CPU 上验证 AMP 关闭后的完整训练步的最小验证：使用完整构造输入核对CPU 上验证 AMP 关闭后的完整训练步，输出只表示本例计算。
变量解释
| 变量 | 含义 |
|---|---|
N | batch size |
D/H | 特征或隐藏维度 |
完整数值示例
CPU 上验证 AMP 关闭后的完整训练步
已知条件
- 单权重线性层初值 1
- x=2，目标=6，学习率 0.1
- autocast 与 GradScaler 均显式关闭
计算步骤
- 前向平方损失为 16。
- scaler 在关闭状态下透传操作。
- 更新后权重为 2.6。
计算结果
CPU 回退使用 float32，得到与普通 SGD 相同的一步更新。
处理前后对比
- 本页 CPU 代码验证回退流程，未测量或宣称 CUDA 性能收益。
CPU 上验证 AMP 关闭后的完整训练步
示例故意不启用不受当前硬件支持的 CUDA 算子。实际切换到 CUDA 时，应把模型与输入迁到同一设备，再按 float16 或 bfloat16 配方启用 autocast 与缩放。
衡量 AMP 应在相同任务预算下同时观察数值稳定性、显存和耗时。单次小矩阵计算的快慢无法代表大网络吞吐，也不能把训练能运行当成精度没有影响。
计算与实现步骤
- 前向平方损失为 16。
- scaler 在关闭状态下透传操作。
- 更新后权重为 2.6。
代码实现
示例 1
CPU 上验证 AMP 关闭后的完整训练步的最小验证
example_01.py用途：使用完整构造输入核对CPU 上验证 AMP 关闭后的完整训练步，输出只表示本例计算。
import torch
from torch import nn
model = nn.Linear(1, 1, bias=False)
with torch.no_grad():
model.weight.fill_(1.)
opt = torch.optim.SGD(model.parameters(), lr=0.1)
scaler = torch.amp.GradScaler("cpu", enabled=False)
opt.zero_grad(set_to_none=True)
with torch.amp.autocast("cpu", enabled=False):
loss = (model(torch.tensor([[2.]])) - 6).square().mean()
scaler.scale(loss).backward()
scaler.step(opt)
scaler.update()
print(f"dtype={model.weight.dtype} loss={loss.item():.1f} weight={model.weight.item():.1f}")代码解析
解析始终位于完整代码下方，并按实际代码段逐项对应。
输入数据与任务
使用完整构造输入核对CPU 上验证 AMP 关闭后的完整训练步，输出只表示本例计算。
model- 组合预处理和估计器的模型对象。
scaler- 保存训练集缩放参数的转换器。
mean- 训练特征逐列均值。
loss- 当前预测与目标之间的损失值。
Step 1 · 1–14 行
导入当前步骤需要的数值计算、预处理、模型或评价工具。依赖集中写在代码开头，便于复现。
import torch
from torch import nn
model = nn.Linear(1, 1, bias=False)
with torch.no_grad():
model.weight.fill_(1.)
opt = torch.optim.SGD(model.parameters(), lr=0.1)
scaler = torch.amp.GradScaler("cpu", enabled=False)
opt.zero_grad(set_to_none=True)
with torch.amp.autocast("cpu", enabled=False):
loss = (model(torch.tensor([[2.]])) - 6).square().mean()
scaler.scale(loss).backward()
scaler.step(opt)
scaler.update()
print(f"dtype={model.weight.dtype} loss={loss.item():.1f} weight={model.weight.item():.1f}")代码与数学原理
减去训练均值并除以训练标准差，对应代码中的 mean/std 或 StandardScaler。
平方误差代码对应均方损失。
预期输出或运行结果
dtype=torch.float32 loss=16.0 weight=2.6
常见错误 · 4 条
- 不要在测试集上调用 fit 或 fit_transform；测试集只能复用训练参数。
- 整个模型强制 float16
- 裁剪前不 unscale
- NaN 不关闭 AMP 定位
实际应用
- CUDA 训练的显存与吞吐优化
- 通过关闭 AMP 定位数值异常
常见错误
- 整个模型强制 float16
- 裁剪前不 unscale
- NaN 不关闭 AMP 定位
常见问题
AMP 时可以先裁剪再 unscale 吗？
不应这样做。梯度仍乘着动态缩放因子，裁剪阈值就不再对应真实梯度尺度，应先 scaler.unscale_(optimizer)。
这个示例能说明什么边界？
本页 CPU 代码验证回退流程，未测量或宣称 CUDA 性能收益。
推荐学习资料
参考库不会生成虚假资源或无效外部链接。