【Bug已解决】Out of memory after some epochs while training RT-DETR v2 model 解决方案一、现象长什么样用 RT-DETR v2 训练目标检测模型时前几个 epoch 显存占用平稳、速度正常但跑到第 N 个 epochN 不固定可能 3、5、8突然抛显存不足CUDA out of memory. Tried to allocate 2.00 GiB (GPU 0; 23.70 GiB total capacity; already allocated 22.10 GiB; 1.20 GiB free; 21.80 GiB reserved in total by PyTorch)更隐蔽的版本不报 OOM但每个 epoch 末尾torch.cuda.memory_allocated()都比上个 epoch 高一点像漏水一样缓慢上涨直到某次empty_cache也救不回来。重启训练后前几个 epoch 又正常过一阵子再次复现——典型的「累积型显存泄漏」。这类问题的迷惑性在于单步看不出问题内存是跨步、跨 epoch 慢慢涨的。很多人以为是 batch 里混入了超大图、或者数据增强偶尔产出巨图但把数据查遍也没找到异常样本。二、背景RT-DETR v2 是实时端到端检测器训练循环大致是for epoch in range(epochs): for images, targets in loader: outputs model(images) loss criterion(outputs, targets) loss.backward() optimizer.step() optimizer.zero_grad()「几个 epoch 后 OOM」几乎不可能是优化器状态膨胀Adam 的动量/方差是固定大小的也不太可能是模型本身变大。真正的元凶通常是某处把张量偷偷留在了 Python 对象里导致它的 autograd 计算图一直不被释放。最常见的两个藏匿点日志列表无脑 append 张量为了画 loss 曲线很多人写loss_history.append(loss)。但loss带着完整的反向图append 进普通 list 后这个图就被 list 引用住了永远等不到释放。每跑一个 batch 就多存一张图跨 epoch 累积显存稳步上涨。这正是「几个 epoch 后 OOM」的头号原因。RT-DETR v2 的特征缓存没清RT-DETR 的 decoder 为了加速会在memory/feat上做多尺度缓存部分实现会把 backbone 输出的多尺度特征存进模块属性供同 batch 的多次 forward 复用。如果缓存用「持久属性」而不是「每步重建的局部变量」实现且 forward 之间没清缓存的特征图动辄几百 MB就会一层层叠起来。下面用最小代码复现第一种最普遍、也最致命的情形。三、根因根因一句话训练循环把「带 autograd 图的张量」长期持有在 Python 容器list / dict / 模块属性里计算图无法被 GC 回收显存随 batch、随 epoch 单调累积最终 OOM。细分带图张量被引用loss_history.append(loss)持有loss及其整个反向图feat_cache backbone_feats持有大特征图。图不被释放PyTorch 的显存释放依赖「无引用即回收」。只要 list/dict/属性还指着它图就活着显存就占着。跨 epoch 累积单步泄漏很小几十 MB但一个 epoch 几百步、几个 epoch 下来就是几 GB超过显存上限就 OOM。不是 batch 太大、不是模型太大是「引用泄漏」。四、最小可运行复现下面不依赖真实检测器用一个小循环演示「append 带图张量」如何吃光显存import torch import torch.nn as nn model nn.Linear(1024, 1024).cuda() opt torch.optim.SGD(model.parameters(), lr1e-3) leak_list [] # 故意泄漏存带图的 loss safe_list [] # 正确做法只存标量 def train(steps, leakFalse): for i in range(steps): x torch.randn(64, 1024, devicecuda) loss model(x).pow(2).mean() loss.backward(); opt.step(); opt.zero_grad() if leak: leak_list.append(loss) # 错误持有整张图 else: safe_list.append(loss.item()) # 正确只存 Python 标量 if i % 50 0: alloc torch.cuda.memory_allocated() / 1e6 print(fstep {i}: 已分配 {alloc:.0f} MB, list 长度 {len(leak_list) if leak else len(safe_list)}) print( 泄漏版本 ) train(200, leakTrue) print( 安全版本 ) train(200, leakFalse)跑泄漏版本你会看到torch.cuda.memory_allocated()随 step 单调上升每步多存一张图而安全版本基本持平。把 step 放大到几千、batch 放大就是「几个 epoch 后 OOM」的精确复现。五、解决方案第一层最小直接修复最小修复任何要长期保存的张量先.detach()再决定存不存只关心数值就.item()转成 Python 标量特征缓存每步清掉。import torch loss_history [] # 仅用于画曲线存标量即可 feat_cache {} # RT-DETR 多尺度特征缓存每步清空 for epoch in range(epochs): feat_cache.clear() # 关键每个 epoch 开头清缓存 for images, targets in loader: outputs model(images) # —— RT-DETR v2 特征缓存用局部变量不要持久属性 —— # 如果模型内部用 self.feat_cache 缓存改为每步重建 # self.feat_cache compute_feats(...) 而不是 self.feat_cache.append(...) # 或在 forward 末尾清空model.clear_feat_cache() loss criterion(outputs, targets) loss.backward() optimizer.step() optimizer.zero_grad() # 正确记录只存标量绝不留图 loss_history.append(loss.detach().item()) # 或 loss.item() # 若必须保留张量如做梯度累积统计至少移回 CPU 且 detach # some_stats.append(loss.detach().cpu())具体针对 RT-DETR v2 的特征缓存在模型侧加一个清缓存钩子class RTDETRv2Wrapper(nn.Module): def __init__(self, model): super().__init__() self.model model self._feat_cache None def clear_feat_cache(self): self._feat_cache None # 释放对大特征图的引用 def forward(self, images, cache_featsFalse): if cache_feats and self._feat_cache is None: self._feat_cache self.model.backbone(images) # 仅在同 batch 复用 feats self._feat_cache if cache_feats else self.model.backbone(images) return self.model.decoder(feats) # 训练循环每步后清 # out wrapper(images, cache_featsTrue); ... # wrapper.clear_feat_cache()要点loss.detach().item()切断图只留 Python floatlist 再大也不占显存。特征缓存用「局部变量 每步清空」不让大特征图被模块属性长期引用。model.clear_feat_cache()在 forward 之后调用断开引用让 GC 回收。这一步单独就能让显存跨 epoch 持平。六、解决方案第二层结构性改进第一层是「在循环里修两处」。但训练脚本散落多处append、缓存、EMA、指标累积最容易漏。更好的做法是用一个单一的内存看守对象统一接管「哪些张量能留、留多久、留什么形态」。from dataclasses import dataclass, field from typing import List, Optional import torch dataclass class TrainMemoryGuard: 训练期显存泄漏的单一看守。 # 允许保留的历史长度上限超过就丢最旧的防御极端累积 max_history: int 1000 # 是否把保留的张量移回 CPU offload_to_cpu: bool True # 峰值显存告警阈值字节0 表示不告警 warn_bytes: int 0 _losses: List[float] field(default_factorylist, reprFalse, initFalse) _peak: int field(default0, reprFalse, initFalse) def record_loss(self, loss_tensor: torch.Tensor): # 永远只存标量绝不持有图 self._losses.append(float(loss_tensor.detach().cpu().item())) if len(self._losses) self.max_history: self._losses.pop(0) def hold_tensor(self, tensor: torch.Tensor, name: str) - Optional[torch.Tensor]: 需要跨步持有张量时统一 detach 可选 offload。 t tensor.detach() if self.offload_to_cpu: t t.cpu() return t def check(self): if not torch.cuda.is_available(): return alloc torch.cuda.memory_allocated() self._peak max(self._peak, alloc) if self.warn_bytes and alloc self.warn_bytes: print(f[TrainMemoryGuard] 显存偏高: {alloc/1e6:.0f} MB) def reset_peak(self): self._peak 0 property def peak_bytes(self): return self._peak property def loss_history(self): return self._losses # 用法 guard TrainMemoryGuard(max_history500, warn_bytes22 * 1024**3) for epoch in range(epochs): for images, targets in loader: out model(images) # RT-DETR 缓存若需跨子步持有用 guard.hold_tensor loss criterion(out, targets) loss.backward(); optimizer.step(); optimizer.zero_grad() guard.record_loss(loss) # 安全记录 guard.check() # 监控 guard.reset_peak()结构收益统一看守所有「记录 loss / 持有张量」都过TrainMemoryGuard不会再有人手滑append(loss)。防御累积max_history兜底即便忘记清也不会无限涨。可观测check()peak_bytes让 CI / 日志能画出显存曲线泄漏一眼可见。七、解决方案第三层断言 / CI 守护写 pytest 守两条铁律(1) 记录 loss 后该 loss 张量的图不再被持有显存不涨(2) 跨步持有张量必须用 detachoffload。import torch import pytest from your_lib import TrainMemoryGuard def test_record_loss_does_not_retain_graph(): if not torch.cuda.is_available(): pytest.skip(需 cuda 验证显存) guard TrainMemoryGuard() model torch.nn.Linear(512, 512).cuda() opt torch.optim.SGD(model.parameters(), lr1e-3) before torch.cuda.memory_allocated() for _ in range(200): x torch.randn(32, 512, devicecuda) loss model(x).pow(2).mean() loss.backward(); opt.step(); opt.zero_grad() guard.record_loss(loss) # 只存标量 after torch.cuda.memory_allocated() # 200 步后显存不应显著增长容忍 50MB 抖动 assert after - before 50 * 1024**2, f显存泄漏: {(after-before)/1e6:.0f} MB def test_hold_tensor_detaches_and_offloads(): guard TrainMemoryGuard(offload_to_cpuTrue) t torch.randn(4, 4, requires_gradTrue) held guard.hold_tensor(t, demo) assert held.requires_grad is False, 持有张量必须 detach if torch.cuda.is_available(): assert held.device.type cpu, 应 offload 到 CPU def test_max_history_caps_growth(): guard TrainMemoryGuard(max_history3) for i in range(100): # 用假张量模拟record_loss 只存标量长度受控 guard.record_loss(torch.tensor(float(i))) assert len(guard.loss_history) 3, 历史应被截断CI 常驻跑这三条后任何「手滑 append 带图张量」「缓存不清」的回归都会立刻爆红。八、排查清单训练「几个 epoch 后 OOM」时按顺序查先确认是不是数据问题打印每步images.shape看有无偶发超大图。没有就怀疑泄漏。全局搜append(loss)/append(output)/history.append(...)确认 append 的是.item()标量还是张量。搜模型里self.xxx_cache 、self.feat 这类持久属性确认 forward 之间有清空。在训练循环里周期性打印torch.cuda.memory_allocated()画成曲线——单调上升就是泄漏的铁证。确认loss.backward()后有optimizer.zero_grad()且没把loss存进任何容器。EMA 模型如EMAModel会多一份参数副本确认它不会在每个 epoch 重复.cuda()新建。若用torch.cuda.amp/GradScaler确认scaler.update()每步都调且没把scaled_loss存起来。九、小结RT-DETR v2 训练「几个 epoch 后 OOM」根子是训练循环把带 autograd 图的张量loss、特征缓存长期持有在 list / 模块属性里计算图无法回收显存随 batch、随 epoch 单调累积。修复三层次第一层记录 loss 只用loss.detach().item()、特征缓存每步清空第二层用TrainMemoryGuarddataclass 统一接管「什么能留、留多久、留什么形态」并加峰值监控第三层用 pytest 守「记录 loss 不涨显存」「持有张量必须 detachoffload」「历史长度受控」。工程启示训练/推理里任何「跨步持有张量」的地方都默认先.detach()只关心数值就.item()要留大张量就detach().cpu()并设上限。显存泄漏从来不是「突然发生」而是「每步漏一点、跨 epoch 攒爆」所以监控memory_allocated曲线比等 OOM 更有用。