autocast()是智能档位切换器,根据运算敏感度自动选择fp16或fp32;需包裹前向+loss全过程,输入保持float,配合.eval()和.no_grad()正确使用。

autocast()不是开关,是“智能档位切换器”
torch.cuda.amp.autocast 不是把整个模型一刀切进 FP16 的暴力转换工具。它本质是一个上下文管理器,会根据当前运算的数值敏感度,自动决定该用 float16 还是 float32:卷积、GEMM、ReLU?放心走 FP16;LayerNorm、Softmax、loss 计算?默认保 FP32。这种“该省则省、该稳则稳”的策略,才是它不掉点又提速的关键。
- 它只对 CUDA tensor 和支持 AMP 的算子生效,CPU 模型或自定义 C++ 扩展不在此列
- 不需要改模型结构,也不影响
model.state_dict()中权重的原始 dtype(加载仍是 FP32) - 如果模型里手动写了
.float()或.half()强制转换,反而可能破坏 autocast 的判断逻辑
三行代码就能启用,但位置错一格就失效
启用 autocast 最常见的失败,不是不会写,而是没包对范围。它必须包裹「前向计算 + loss 计算」全过程,且不能漏掉任何中间 tensor 的生成环节。
with torch.cuda.amp.autocast():
pred = model(img) # ✅ 在 autocast 内
loss = criterion(pred, gt) # ✅ loss 计算也必须在内
- ❌ 错误示范:
pred = model(img)在外,loss = criterion(...)在内 → pred 是 FP32,loss 反而用 FP16 算,可能报错或溢出 - ❌ 更隐蔽的错:
img = img.half()提前转了半精度 → autocast 失效,后续全按 FP16 跑,BatchNorm 统计量崩坏 - ✅ 正确姿势:输入保持
img.float().cuda(),让 autocast 自主决策每一步
推理时不需要 GradScaler,但显存释放有陷阱
GradScaler 是训练专属组件,用于梯度缩放防下溢;纯推理场景完全不用它。但很多人开了 autocast 后发现显存没降——问题往往出在 tensor 生命周期上。
- 推理中若保留中间结果(比如缓存
feature_map供后续可视化),这些 tensor 会被 autocast “染色”为 FP16,长期驻留显存 -
torch.no_grad()必须和autocast()同级嵌套,顺序不能颠倒:with torch.no_grad(), torch.cuda.amp.autocast(): out = model(x) - 若用了
torch.inference_mode()(PyTorch ≥ 2.0),它本身已隐含类似 autocast 的优化,再套一层 autocast 可能冗余甚至冲突
模型加载和 eval() 顺序不能反
FP16 推理效果好不好,一半看 autocast,另一半看模型是否真正“准备好”。常见掉坑点是:模型还在 train 模式、BN 层没冻结、或权重没进 GPU 就启用了 autocast。
- 必须先调用
model.cuda().eval(),再做任何推理相关操作 -
.eval()不只是关 dropout,更关键的是让 BatchNorm 切换到使用 running stats 模式;若漏掉,FP16 下 BN 的 FP32 统计量和 FP16 输入混合运算,极易数值漂移 - 不要对模型整体调
.half(),尤其当模型含nn.BatchNorm2d或nn.LayerNorm时,会直接破坏其内部 FP32 参数
autocast 的边界很清晰:它只管计算过程,不管模型状态、不碰权重加载、也不负责显存清理。真正卡住性能的,往往是那几行看似无关紧要的 .eval()、.no_grad() 和 tensor 类型控制。











