
最近我把 YOLO26 的训练流程从 FP32 切到混合精度训练AMP之后最直观的感受是终于不用再天天盯着 nvidia-smi 算显存余额了。同样的 batch size同样的 640 输入分辨率显存占用直接降了接近四成训练速度还反超了一截。这篇文章把我这次实战改造的思路、代码改动和踩坑记录完整写出来给正在被显存卡脖子的同学一个可以直接抄作业的参考方案。AMP 是什么简单说它不是在训练质量上糊弄而是把计算图和中间结果里那些不需要高精度的部分切成半精度FP16/BF16该保留高精度的地方继续用 FP32。YOLO 这类检测模型对数值精度相当宽容正因如此AMP 在 YOLO26 上能省下这么多显存而且几乎不掉点。这篇内容适合谁看用 8G/12G/16G 显存显卡训练 YOLO26、想在不降 batch size 的情况下塞进更大输入尺寸、以及被 OOM 反复折磨的炼丹师。我会从显存到底花在哪讲起再逐步拆解混合精度的原理、实操接入方式、实测数据和避坑经验全程附可复现代码。1. 显存瓶颈从哪来先算清这笔账1.1 训练时显存开销的大头在哪很多人误以为训练模型最吃显存的是模型参数本身其实不是。以 YOLO26 为例假设模型参数量在 4000 万到 5000 万这个量级FP32 的模型权重大约只占 160 到 200MB放在一张 12G 显存的卡上根本谈不上压力。真正把显存写满的是另外三块内容前向传播过程中为反向传播保存的中间激活值feature map、优化器保存的动量与方差状态、以及每个 batch 里的图像和标注数据。激活值是最容易被忽略的大头。YOLO26 在 640x640 输入下会产生多个尺度的特征图网络越深、通道数越大每一层保存下来的激活张量就越多。Batch size 为 16 时光这些中间特征累积起来就能轻松吃掉 4 到 6GB 显存。反向传播要计算梯度就需要知道每一层的输入和输出这些张量没法省省了就只能用重计算去换时间。优化器状态同样不容小觑。用 Adam 系优化器训练时每个参数不仅要存本身还得额外存一阶动量 m 和二阶动量 v。一个参数在 FP32 下需要 4 字节权重、4 字节梯度、再加 8 字节动量状态也就是 16 字节。4000 万参数就是 640MB如果是 8000 万参数的模型这一项直接超过 1.2GB。很多二手卡看起来显存不小真跑起来才发现被参数和状态吃得干干净净。1.2 FP32 为什么浪费了将近一半的存储FP32 用 32 位存储一个浮点数FP16 用 16 位存储比特位减半理论上显存占用也能跟着减半。但这个减法不是简单切一刀因为两种精度能表示的数值范围完全不同。FP32 的 32 位里1 位给符号8 位给指数23 位给尾数FP16 则是 1 位符号、5 位指数、10 位尾数。指数位减少意味着动态范围变窄尾数位减少意味着同一数量级下的相对精度变低。精度类型占用字节指数位尾数位最小正正规数适用场景FP324823约 1.18e-38主权重、优化器状态FP162510约 6.10e-5前向激活、卷积计算BF16287约 1.18e-38大模型梯度、极端动态范围从表格能看出FP16 的动态范围小了很多尾数精度也低了不少但最关键的其实是后面两行FP16 最小的正规数是 6.10e-5 左右而 FP32 可以小到 1.18e-38。深度学习中很多梯度的绝对值非常小一旦落到 FP16 的表示范围以下就会直接变成 0这就是所谓下溢。如果一路在前向计算时用 FP16反向传播时梯度直接消失模型自然训不动。1.3 为什么不直接用 FP16 训练而要混合既然 FP16 有下溢问题为什么不把所有东西都换成半精度答案就是标题里那两个字混合。AMP 的核心思路很直白——权重、优化器状态、梯度累积这些对精度敏感的部分保留 FP32作为“主权重”存在而前向传播中的卷积、批归一化之后的激活、以及大量中间计算结果则用 FP16 计算充分发挥硬件加速能力。主权重用 FP32 还有一个好处每次更新时即使梯度的绝对值非常小在 FP32 累加后也不会丢精度。等累积到一定程度再与学习率相乘数值依然稳定。这条设计路线是英伟达在 Volta 架构之后推广起来的标准玩法现在 PyTorch 里的 autocast 和 GradScaler 就是为这个流程量身定做的。理解了为什么混合后面接入代码时就不容易配错。2. 准备工作搭好混合精度训练的基础环境2.1 硬件要求不是所有 GPU 都适合开 AMP混合精度训练对显卡有最低门槛。建议至少使用支持 FP16 硬件加速的 GPU也就是计算能力 6.1 以上的显卡比如 GTX 10 系列、RTX 20/30/40/50 系列以及 T4、V100、A100 这类专业卡。只要 GPU 支持 FP16 指令AMP 就能跑训练代码本身不会报错但老架构没有 Tensor CoreFP16 计算还是走普通 CUDA 核心速度上几乎得不到提升甚至可能因为精度转换开销而更慢。从实际体验看RTX 3060 12G、RTX 3070 8G 这类卡跑 YOLO26 AMP 效果非常明显尤其是显存只有 8G 的平台FP32 模式下 batch size 稍微调大就 OOM切到 AMP 之后往往能多塞一倍的 batch 进去。如果是 GTX 1050Ti 这种 4G 老卡建议直接考虑降低输入分辨率别指望 AMP 能创造奇迹。你要是在训练日志里看到 AMP 开启后速度没有变化先查一下显卡的计算能力多半是硬件不支持 Tensor Core。2.2 软件栈与版本核对PyTorch 从 1.6 开始原生支持自动混合精度现在推荐直接用 PyTorch 2.xAPI 更稳定而且 AMP 相关接口已经集中在 torch.amp 命名空间下。我在这次 YOLO26 项目里用的是 PyTorch 2.1.2 CUDA 11.8ultralytics 包用的是最新版本这套组合跑下来没有出现兼容性毛病。下面两条命令可以先跑一下确认环境没问题python -c import torch; print(torch.__version__, torch.version.cuda) python -c import torch; print(torch.cuda.get_device_capability(0))注意第二条命令输出的计算能力元组第一个数字如果是 6 或 7说明支持 FP16如果是 7.5 以上Tensor Core 会明显提速。版本差异还有一个坑老项目里常见的from torch.cuda.amp import autocast, GradScaler写法在 PyTorch 2.x 中仍然能用但官方已经推荐迁移到from torch.amp import autocast, GradScaler新项目尽量用新接口避免后续升级时收到 DeprecationWarning。2.3 数据管线准备好防止“脏数据”干扰半精度很多新手切到 AMP 后第一次跑就遇到 NaN第一反应是 AMP 的问题其实大概率是数据管线本来就有隐患。半精度会把问题放大如果数据加载时某些图像像素出现 NaN 或 Inf比如解码损坏的 JPEG 图片、标注框坐标越界FP32 下可能靠隐式容错勉强过去FP16 下就彻底爆掉。建议在训练前做一次数据体检。用脚本遍历训练集把张量里的 NaN、Inf 值筛查一遍顺便检查图像尺寸是否统一切到目标分辨率。YOLO26 的官方数据加载流程已经做了不少防护但如果你是自己写 Dataset一定记得在 collate 之后做一次 clamp。另外图像归一化最好保持在前向之前的全图上做不要在装进 GPU 之后再用 FP16 去算归一化那样容易产生不必要的精度损失。3. 实操为 YOLO26 训练接入混合精度3.1 基于 ultralytics 的一行开关如果你的 YOLO26 是用 ultralytics 这套框架训练的那混合精度已经被内置封装好了根本不用手写 AMP 逻辑。只需要在训练参数里把 amp 设为 Truefrom ultralytics import YOLO model YOLO(yolo26.yaml) model.train( datacoco.yaml, epochs100, imgsz640, batch16, ampTrue, device0, cacheTrue, )细心的同学可能会发现ultralytics 的 amp 参数默认就是 True也就是说你不传它也会打开。但我会建议显式写上一是让实验记录里的配置一目了然二是方便在做 FP32 对照组时改成ampFalse。框架内部会自行处理 GradScaler、autocast 以及验证阶段的精度回退对大多数业务场景来说足够可靠。框架内置方案的局限在于它对自定义模型结构或自定义损失函数不够灵活。如果你在 YOLO26 上改了检测头、加了辅助分支或者用了完全自定义的损失计算方式那还是建议把训练循环拆出来用手写 AMP 的方式把关键位置控制在自己手里。3.2 手写训练循环的 AMP 接入点自定义训练循环时AMP 的接入点其实就三个前向计算外面套一层 autocast、反向传播之前用 GradScaler 对 loss 做缩放、优化器更新之后调用 scaler.update()。完整代码模板如下import torch from torch.amp import autocast, GradScaler device cuda scaler GradScaler(cuda) model.train() optimizer.zero_grad() for images, targets in dataloader: images images.to(device, non_blockingTrue) targets [t.to(device, non_blockingTrue) for t in targets] optimizer.zero_grad() with autocast(cuda, dtypetorch.float16): outputs model(images) loss compute_loss(outputs, targets) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()这里的 compute_loss 指的是你自己的损失函数YOLO26 的损失包括分类损失、边界框回归损失和置信度损失它们相加后仍然是普通张量。把 loss 计算也放进 autocast 块里是为了让损失函数内部的约简和中间结果同样走半精度从而进一步压低显存占用。如果你发现某些损失项数值波动剧烈也可以在 autocast 块外面单独用 FP32 计算这一项再把最终 loss 传回来这在多任务损失里是常见操作。还有个很容易被忽略的点scaler.step内部做了对梯度的反缩放它会检查梯度是否存在 NaN/Inf如果检查到异常会跳过这次优化器更新。也就是说你在判断“是否真正更新参数”时不能只看optimizer.step()而是要看scaler.step(optimizer)有没有实际执行。想精确控制的话可以参考源码用scaler.unscale_(optimizer)先反缩放再手动 clip。3.3 损失缩放Loss Scaling的作用与默认参数GradScaler 的作用是防止 FP16 梯度下溢。训练时它会把 loss 放大一定倍数反向传播得到的梯度也会被等比例放大这样那些绝对值很小的梯度就能落在 FP16 可表示的范围内。等到真正更新权重时再对梯度做反缩放相当于把放大镜撤掉恢复真实梯度值。PyTorch 里 GradScaler 的默认参数是init_scale2^16、growth_factor2.0、backoff_factor0.5、growth_interval2000。启动时梯度先缩放 65536 倍如果连续 2000 步没有出现 inf/NaN缩放因子会乘 2 继续变大一旦某一步发现溢出scale 会立即减半以此自动寻找当前训练阶段最合适的缩放水平。实际训练中你几乎不需要手动调整这些参数除非遇到梯度爆炸特别频繁的情况。此时可以先把init_scale调小到2^14并关闭动态增长改用固定 scale。我踩过一次坑模型前向里有一段自定义的三角函数运算在 FP16 下输出范围很容易超导致每几十步就出现一次loss nan直到把该层改成 FP32 计算才消停。这说明损失缩放解决不了所有精度问题它只是保证下溢不被无视真正的病根还得靠算子和层设计来治。4. 实测数据显存真的降低 40% 吗4.1 实验对照设计为了保证数据有说服力我专门在同一台机器上跑了 FP32 和 AMP 两组实验。硬件是 RTX 3080 10G数据集选了 COCO 的一个子集总共约 2 万张图训练 30 个 epoch输入分辨率 640x640batch size 设为 16。唯一变化的变量就是 amp 开关其他包括随机种子、优化器、学习率策略全部保持一致。对比的指标有三个训练全程峰值显存占用、每轮平均耗时、以及验证集上的 mAP50。峰值显存用 nvidia-smi 里的 max 值读取训练耗时取最后 10 个 epoch 的平均值避免启动阶段和缓存预热带来的误差。这里说明一下我只测了单卡训练模式没有开 DDP因为大多数人遇到显存瓶颈都是单卡场景保持单卡对照更贴合实际。4.2 显存占用、训练速度和精度的三维对比指标FP32 训练AMP 训练变化峰值显存占用10.4GB6.2GB降低约 40%每轮训练耗时约 152s约 118s提速约 22%mAP500.6340.631基本持平是否出现 OOMbatch 16 勉强能跑余量充足-FP32 这边峰值 10.4GB 已经非常接近 10G 卡的物理上限中间只要多开一个可视化工具都可能触发 OOM。AMP 这边直接把峰值压到 6.2GB节省幅度在 40% 左右核心贡献来自激活值全部变成 FP16一批图像和中间特征图省下来的内存非常可观。训练速度也快了 22%这部分是 Tensor Core 加速卷积计算的功劳。精度方面mAP50 从 0.634 降到 0.631下降 0.003。这个差距在检测任务里属于正常噪声范围同一个随机种子跑两遍 FP32 可能都会有这么大波动。如果你在自家数据集上看到精度下降超过 0.01就需要排查是不是某些层在 FP16 下不稳定或者数据里存在异常值。4.3 为什么你复现出来的数字不一定是整 40%我要强调的是40% 这个数字不是放之四海皆准的结论。显存节省比例取决于模型宽度、输入分辨率、batch size 以及优化器类型。如果 YOLO26 用的模型规模较小参数和优化器状态占比上升激活值占比下降AMP 能省的显存比例就会跌破 35%反之如果你把输入分辨率从 640 提到 1280激活值占据主导地位节省比例很容易超过 45%。batch size 的影响也很微妙。batch 越小显存里固定开销参数、优化器状态占比越大AMP 的节省比例越低batch 越大激活值占比越高AMP 的收益越明显。所以如果你在小 batch 下复现出“怎么只有 20%”不用惊讶这是正常的量级效应。实际操作中用 AMP 的最大价值不在省那百分之几十而在于把原本放不下的 batch 放进显存从而避免降低 batch size 带来的收敛质量损失。5. 避坑指南混合精度训练的常见问题5.1 训练中遇到 NaN/Inf 怎么办AMP 训练中遇到 NaN 时先别急着关混合精度按下面的层次排查。先用固定随机种子复现问题排除数据偶然性然后在训练循环里打印 loss 的数值和 GradScaler 的当前 scale观察 loss 是在哪一步突然变成 NaN 的再用torch.isnan(torch.max(out))排查前向输出的每个关键层。常见原因有四类第一是损失函数里出现了 log 0 或 sqrt 负数FP16 下更容易触发尤其是自定义损失时第二是学习率过高导致梯度更新幅度超出 FP16 的表示范围可以尝试降低学习率或用 warmup第三是训练数据里有脏数据检查数据集是否有损坏图像第四是检测头里用了精度敏感的算子比如某些自定义 attention 在 FP16 下不稳定解决办法是把这一层放在 autocast 块外面或者强制用 FP32 计算。我自己的经验是YOLO26 主体结构跑 AMP 通常很稳出问题的几乎都是后加的模块。所以排查时优先检查你的自定义分支而不是怀疑模型主干。5.2 BatchNorm 与 AMP两个易踩的隐蔽坑BatchNorm 是 AMP 场景里最容易埋雷的结构。PyTorch 的 autocast 机制默认对 BatchNorm 的输入输出做了特殊处理会让它内部以 FP32 累积所以正常情况下你不用额外操作。但如果你在网络前向里手动把输入images.half()或者用model.half()把整个模型切到半精度BatchNorm 就会在 FP16 下工作批内均值和方差的计算精度会显著下降长时间训练会出现验证指标退化。第二个坑是自定义的归一化层。如果你把 YOLO26 里的 BN 换成了其他归一化方式比如 InstanceNorm 或 GroupNorm这些层在 autocast 下不一定自动走 FP32需要你自己确认。我的建议是在训练循环内不要对输入张量做任何显式的.half()转换只依赖 autocast 的上下文管理让框架决定每个算子的精度。5.3 显存降下来了还能怎么继续压榨AMP 只能解决一部分显存问题如果省下来的显存还是不够用我建议按顺序叠加下面几种手段。第一是梯度检查点gradient checkpointing它用重计算换显存能把激活值占用再降低一半左右但训练速度会多花 20% 到 30%建议在 AMP 的基础上按需开启。第二是适当缩小输入分辨率YOLO26 对分辨率比较敏感但你可以用随机多尺度训练让模型适应不同尺寸同时验证时只测最高分辨率显存曲线会平滑很多。第三是开 DDP 多卡训练。AMP 加上 DDP 的组合在显存效率上几乎是线性扩展原来单卡 10G 勉强跑 batch 16双卡只用 5G 就能跑同样 batch还可以把 batch 翻倍。最后如果目标平台是低显存设备训练完成后的模型转换也很重要把最佳权重导出成 FP16 的 ONNX再用 NCNN 转换工具生成 bin 和 param就能在几百 MB 显存的嵌入式设备上做推理。这一套组合拳下来才能真正把显存瓶颈从训练端到部署端都解决掉。我在这次改造里最深的体会是AMP 不是靠“降低精度”偷工减料而是把精度资源分配到了真正需要它的地方。它不需要你有非常深的底层知识但需要你理解哪些地方能省、哪些地方不能省。如果你现在正被显存卡住先不用忙着换卡把 YOLO26 的混合精度训练打开实测一轮很可能会有意外惊喜。