YOLO26 低显存训练优化：AMP 混合精度实战，显存直降 40%

最近的活是拿 YOLO26 训一批工地头盔检测的数据显卡只有 12G 的 RTX 3060。第一轮 FP32 跑起来batch 开到 16 直接卡死在 OOM把 imgsz 降到 512 才勉强活过来。后来认真把混合精度训练AMP从头到尾调了一遍同一个模型、同样 640 分辨率峰值显存从 10.4GB 掉到 6.2GB降幅正好 40% 出头训练吞吐还涨了 20% 以上。这篇文章不是标题党式的“试一试混合精度”而是把 YOLO26 在低显存环境下跑训练的关键细节全部摊开讲。如果你跟我一样卡在显存瓶颈正为“模型能跑起来但 batch 开不大”发愁下面的内容可以直接照着抄。我会先说清楚显存到底被谁吃了再解释 AMP 为什么能省这么多最后给出完整配置和踩坑记录。1. 模型还没跑起来显存已经被“四分五裂”了先搞清楚谁在吃显存1.1 训练进程里显存其实分成了好几块独立的开销很多人一看到 OOM 就以为模型权重太大了实际上训练进程里的显存是好几块开销叠加出来的权重往往只占很小一部分。我习惯按下面五块去拆解CUDA context 和 cuDNN workspace这是 PyTorch 初始化后固定占用的“背景内存”8G、12G 卡上差不多 0.4~0.8G省不掉。模型权重和 BN bufferyolo26s 参数量大概 20M 上下FP32 权重也就 80~100MB真的不大。优化器状态用 AdamW 的话每个参数额外要存一份一阶动量 m 和二阶动量 v都要 FP32这部分通常和模型权重差不多大小甚至更大。前向激活值每一层输出在反向传播时都要被保存供计算梯度使用。它跟 batch size、分辨率、特征图尺度强相关是显存峰值里最不可控的大头。梯度缓冲和反向计算图梯度的大小跟激活值同数量级AMP 能砍掉接近一半。用一个生活化的类比模型权重是快递箱里的货激活值就是填充泡沫。货可能不重但泡沫把整个箱子撑满了。训练时真正把显存塞爆的大多数情况下是那一堆“做完前向、等着反向时再拿出来用”的中间张量。1.2 为什么第一条命令就 OOM典型大头是激活值而非权重我在跑 YOLO26 之前也犯过这个误区觉得模型文件才一百多 MB显存 12G 怎么会不够实际以imgsz640, batch16为例yolo26s 的前向激活值能轻松堆到 3~4G 甚至更多如果模型是 m 或 l 版neck 和 head 部分的特征图会把激活值推到 6G 以上。加上优化器状态、CUDA context、DataLoader 预取8G 卡不动是正常的12G 卡也得小心翼翼。所以先记住一个结论训练显存的峰值七八成由 feature map 激活值决定。这决定了 AMP 为什么能起这么大作用也决定了后面说的梯度检查点为什么是第二板斧。2. AMP 到底省的是哪块显存别把“省一半”想简单了2.1 混合精度不是“全切 FP16”而是三项分离一听到混合精度很多人以为就是把模型所有张量都从 FP32 换成 FP16显存直接减半。真这么干训练大概率会 NaN 满天飞。AMP 的思路是三件事分开处理模型参数和优化器状态继续保留 FP32 的 master copy用于更新权重保证数值稳定。前向计算和激活值保存进入 autocast 上下文后卷积、BN、激活函数这些算子自动用 FP16/BF16 计算中间激活值也以 FP16 保存。反向梯度和梯度缩放梯度算出来存成 FP16但为了防止小梯度在 FP16 下直接下溢成 0前向结束以后会把 loss 乘上一个很大的 scale 因子反向完成后在优化器 step 之前再除以这个 scale最后真正更新的是 FP32 权重。简单说AMP 是“用 FP16 省内存和省时间用 FP32 保精度用 loss scaling 防梯度消失”。它不是把所有东西切成一半而是该 FP32 的继续 FP32该 FP16 的才切。2.2 省下来的具体比例为什么是 40%不是 50%理论上激活值和梯度切成 FP16这部分确实能省一半但整体降幅不会是 50%因为 CUDA context、模型权重、优化器状态都没怎么降。我在 RTX 3060 上测到的基线是这样的峰值 10.4G 里激活值和反向图加起来大概 6.5G压成 FP16 应该省出 3.2G理论峰值大约 7G。实际跑出来是 6.2G这说明 PyTorch 的显存分配器在用了 FP16 张量后缓存池碎片更少、预留量也变小了叠加下来才真正到了 40% 左右。如果你的模型比较小、分辨率比较低激活值占比没那么高省幅可能只有 20%反过来模型越大、batch 越大激活值占比越高省幅就越接近 50%。我建议不要只信网上“混合精度能省 40%”这句话用自己基线的激活占比去估这才是最靠谱的。2.3 为什么 YOLO26 这种密集卷积结构特别吃 AMP 红利YOLO26 的 head 部分跨了多个尺度输出其中 stride 8 那一路的特征图分辨率很高neck 里又有不少 C2f 结构这些地方会把前向中间结果大量保留。相比那些轻量级分类模型YOLO 这类检测网络对“激活值显存”更敏感所以 AMP 的收益往往比分类网络更明显。换句话说你不是在省模型的“体重”而是在省训练过程中临时铺开的“包装材料”。3. 动手前先做一次显存体检基线数据比感觉靠谱3.1 环境准备和数据集格式先统一我这边用的 YOLO26 是社区分支里的 26 系模型训练入口和 Ultralytics 保持了一样的命令行习惯。环境建议直接建独立的 conda 环境避免把机器的 PyTorch 版本搞乱conda create -n yolo26 python3.10 -y conda activate yolo26 pip install torch2.3.1 torchvision0.18.1 --index-url https://download.pytorch.org/whl/cu121 pip install ultralytics如果你的 YOLO26 是克隆的独立仓库记得先git clone到本地再在仓库目录里执行pip install -e .这样仓库里自带的模型配置才能被命令行识别。数据集格式不要自己造轮子YOLO 家族从 v5 开始就是一套约定俗成的结构datasets/mydata/ ├── images/ │ ├── train/ │ └── val/ └── labels/ ├── train/ └── val/对应的mydata.yamlpath: datasets/mydata train: images/train val: images/val names: 0: person 1: helmet 2: car注意 label 文件的类别编号必须从 0 开始连续编号否则 YOLO26 训练时矩阵计算会直接报越界这个错和气精度一点关系都没有但是最容易在新手身上浪费半天。3.2 用训练前 5 分钟拿到你的显存基线不要一上来就改 AMP先跑一个不开 AMP 的基线。我建议用 1 epoch 快速探一下峰值脚本里显式加上显存统计比在另一个终端盯nvidia-smi准得多import torch from ultralytics import YOLO torch.cuda.reset_peak_memory_stats(0) model YOLO(yolo26s.yaml) model.train( datamydata.yaml, epochs1, imgsz640, batch16, device0, ampFalse, workers0, ) print(f已用完显存峰值: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB) print(f缓存池预留峰值: {torch.cuda.max_memory_reserved() / 1024**3:.2f} GB)这里有个细节max_memory_allocated是 PyTorch 真正分配给张量的峰值max_memory_reserved是显存分配器从显卡上划走的缓存池大小。nvidia-smi看到的是 reserved 那一档实际有效张量只占 allocated 部分。如果两者差值很大说明存在显存碎片问题后面我会专门讲怎么处理。我自己的基线就是这一步测出来的batch16, imgsz640, yolo26s, ampFalse峰值 10.4G12G 卡已经有点慌8G 卡直接不用想了。先拿到这个数再去谈优化才有参照系。4. 实战FP32 切 AMP 的完整配置与实测结果4.1 两行命令完成切换但配置要显式写YOLO 训练入口在命令行里是这样切换的# 基线不开 AMP yolo train modelyolo26s.yaml datamydata.yaml epochs100 imgsz640 batch16 device0 ampFalse # 正式训练开 AMP yolo train modelyolo26s.yaml datamydata.yaml epochs100 imgsz640 batch16 device0 ampTrue我更推荐在 Python 里跑参数更直观出现问题时也容易加调试逻辑from ultralytics import YOLO model YOLO(yolo26s.yaml) results model.train( datamydata.yaml, epochs100, imgsz640, batch16, device0, ampTrue, optimizerAdamW, lr00.001, lrf0.01, warmup_epochs3, workers4, )有一个坑必须提醒Ultralytics 默认在 CUDA 环境下amp本来就是 True很多人没写这个参数也以为自己在用 AMP。但如果你换了 MPS 或 CPU 设备这个参数会被自动忽略日志里也不会醒目标出。所以要么显式写要么训练一开始就看日志里有没有AMP: True这行别等到跑完了才发现从头到尾都是 FP32。4.2 我的实测结果40% 是怎么来的同一张 RTX 3060 12G同一个数据集yolo26simgsz640batch16训练 80 轮结果如下指标FP32 基线AMP 训练变化峰值显存10.4GB6.2GB降低 40.4%每秒训练帧数42.752.1提升约 22%mAP0.50.9240.928基本持平单卡可开最大 batch1632显存余量翻倍损失曲线偶尔有尖峰稳定依赖 loss scaling更稳需要注意的是这个 40% 是在“相同 batch、相同分辨率”前提下对比的。AMP 省出来的显存如果你立刻拿去开更大的 batch 或更高的分辨率峰值显存会重新涨回来但换到的是更高的吞吐量或精度这才是 AMP 的正确打开方式后面会展开讲。4.3 为什么盯着 mAP 而不是只看显存混合精度训练最诱人的地方是“内存降了速度还快”但前提是精度别崩。我的经验是AMP 改完以后 mAP 和 FP32 相差应该在 ±0.005 以内如果掉了超过 0.03问题通常不在 AMP 本身而在学习率和 warmup。我遇到过一种情况FP32 下用lr00.01能正常收敛切到 AMP 后前几个 epoch 直接 NaN。这不是说 AMP 不能用而是因为 FP16 动态范围小大学习率在第一轮就把 FP32 master weight 冲出了问题。后来我把lr0降到 0.001warmup_epochs调到 3问题消失。所以遇到 AMP 训练不稳先别急着关 AMP先看看你的学习率是不是为了 FP32 调得太激进。5. 省完显存别停在这里batch、分辨率、梯度检查点该怎么组合5.1 显存省下来优先加 batch 而不是只留在原地很多人开了 AMP 后看到显存从 10G 掉到 6G 就觉得完工了。我建议你再多做一步把省出来的显存换成更大的 batch。大的 batch 让 BN 统计值更稳梯度收敛轨迹也更平滑。实测中我做了三组对比方案峰值显存训练速度备注FP32batch1610.4G42.7 imgs/s基线AMPbatch166.2G52.1 imgs/s显存最低AMPbatch3210.1G68.4 imgs/s同等显存吞吐最高AMPbatch32imgsz76811.4G55.3 imgs/s精度收益明显第三组是把省下来的显存换成 batch第四组是换成分辨率。如果你的数据集里有大量小目标优先提imgsz如果数据分布均衡、只是觉得训练慢优先加 batch。不要两个都拉满以免再次 OOM显存管理本质上是取舍。5.2batch-1、imgsz768、cacheram的稳妥组合YOLO26 命令行支持batch-1会自动探测当前显存能开的最大 batch。这个功能在 AMP 前后表现很不一样FP32 下它可能只给到 14 或 16AMP 下它会给到 32 甚至更高。我现在的习惯是先固定imgsz用batch-1探一次上限再手写下限保守 10% 的 batch。AMP 把 GPU 算力释放以后数据加载经常会成为新瓶颈。表现为 GPU 使用率忽高忽低显存下降但训练速度没上来。我通常配合cacheram把数据集预加载到内存再根据 CPU 核心数把workers调到 4~8。注意cacheram吃的是 CPU 内存32G 内存机器上装几个 G 的数据集没问题但如果数据量几十 G就别开这个。5.3 梯度检查点第二板斧但不是免费的如果 AMP 之后你还是想再省一点或者要上更大模型 yolo26m/l可以开梯度检查点gradient checkpointing。它的原理很朴素前向时不保留全部激活值只保留一小部分反向时再用当前层重新算一遍激活值。这样显存能再省不少但代价是额外的前向计算训练速度会掉 20% 左右。如果你的 YOLO26 仓库在 model yaml 里带了gradient_checkpointing参数直接在训练命令行加一行即可。如果没有也可以在自己的模型定义里手动包一层from torch.utils.checkpoint import checkpoint # 将耗显存的 C2f 模块改成 checkpoint 版本 x checkpoint(self.c2f, x, use_reentrantFalse)我的建议是先 AMP再看显存余量决定要不要梯度检查点。AMP 是“收益远大于代价”的一板斧梯度检查点则是“用时间换空间”的第二板斧正常情况下没必要两个一起上训练速度会变得很难看。6. 混合精度翻车现场三个我撞过的坑和最终避坑办法6.1 坑一第一批 loss 直接 NaN经验不足的人第一反应是关掉 AMP前面提到过我在第一次切 AMP 时第三轮 loss 突然变 NaN训练直接崩掉。我当时也本能地想回到 FP32但后面仔细定位发现问题不是混合精度本身而是学习率和 warmup。FP32 能扛住大学习率带来的瞬时梯度波动FP16 的范围更窄扛不住。避坑做法有三件套lr0从原来的 0.01 降到 0.001warmup_epochs3给 GradScaler 的 scale 值一个自适应的过程如果还出错给 train 方法加max_grad_norm10.0限制梯度爆掉的幅度。这三招是在真实训练里验证过有效的。不要一看到 NaN 就把ampTrue改成ampFalse那样等于把显存省下来的路又堵回去了。6.2 坑二显存确实降了但reserved比allocated高一大截AMP 之后张量都变成了半精度数量更多但单个体积更小反而容易把 PyTorch 的缓存分配器切成很多碎片。我遇到过allocated只有 5G但reserved却显示 8G 的情况这就是碎片化导致缓存池一直涨不回收。解决办法在训练前设置环境变量export PYTORCH_CUDA_ALLOC_CONFexpandable_segments:Trueexpandable_segments会让显存分配器按需扩展段而不是从一开始就预留一大块对降低reserved峰值很有帮助。如果你的驱动或者 PyTorch 版本不支持这个选项退一步用max_split_size_mb:128也能缓解。设置完以后重新跑一次第一节里的体检脚本对比一下 reserved 和 allocated 的差值应该会明显缩小。6.3 坑三AMP 提速 22% 以后数据加载变成新瓶颈FP32 的时候 GPU 算得慢DataLoader 的加载时间被计算时间掩盖了看着好像没瓶颈。AMP 把 GPU 算力释放出来以后你会发现nvidia-smi里的 GPU利用率在 60%~100% 之间抖训练日志里的速度忽快忽慢。这通常是 CPU 端的数据增强、图像解码和缩放跟不上了。我的做法是workers8前提是 CPU 核心足够小数据集直接cacheram不要用cacheram的同时再开很大的workers否则 CPU 内存直接爆。还有一个细节容易被忽略AMP 只影响训练阶段的 forward/backward验证阶段通常还是全精度。如果你的验证分辨率设得比训练还高可能会出现“训练不 OOM、验证突然 OOM”的情况。这时候单独把val的imgsz调低或者用fraction0.8控制 PyTorch 最多使用 80% 显存给验证留出余量。我现在的习惯是任何模型在正式开跑之前都先用第一节的脚本做一张显存基线表把 FP32、AMP、AMP更大 batch 三组配置的显存、速度、mAP 全部记下来。这样做的好处是每次训练出问题都能快速判断是模型问题、数据问题还是训练策略问题而不是靠感觉瞎调。混合精度这件事一旦把原理和基线摸透40% 的显存降幅真的只是顺带的结果。