最近在跑一个图像分割的模型,用的DeepLabV3+,backbone是ResNet101。我把batch_size降到2了,输入图片也缩到256x256,但显存还是从开始的2G一直涨到12G,最后OOM。我查了网上说可能是梯度累积、或者变量没detach的问题,但我没开梯度累积,损失函数也是常用的CrossEntropy。想问问大家有没有什么成熟的debug思路?比如用torch.cuda.memory_summary()看哪里泄露,或者有没有工具能可视化每层的显存占用?另外,是不是我模型里有循环或者多次forward导致的?感谢各位大佬!
PyTorch训练时显存一直涨但batch_size已经很小了,咋排查?
全部回复
共 180 条我之前也踩过类似的坑,显存从2G一路涨到12G基本可以排除是batch_size的问题,更像是计算图没释放或者有隐藏的引用在累积。你可以先试试在训练循环里加torch.cuda.empty_cache()看有没有缓解,但注意这只治标不治本,核心还是找源头。torch.cuda.memory_summary()确实值得跑一下,重点看allocated和reserved的差值,如果reserved特别高可能是缓存碎片化,但你说稳定上涨到OOM,更像是每次迭代都有新tensor被保留。另外建议检查一下模型里有没有类似dropout或者batchnorm在eval/train切换时导致的额外forward,DeepLabV3+的ASPP模块里如果有循环遍历不同dilation rate,有时候不小心把中间结果append到list里没清理也会这样。你也可以试试用torch.profiler或者py-spy去抓每个op的显存分配,我上次就是靠这个发现是某个自定义的损失函数里多算了一次log_softmax没detach。对了,确认一下你的优化器是不是只在零梯度时step,有时候梯度累积的隐性bug是写在dataloader里的,比如Sampler重复返回同一批数据。
之前跑分割模型也遇到过这种显存缓慢上涨的情况,最后发现是dataloader的num_workers开太多,每个worker预加载的batch在反向传播后没被及时释放,把workers降到4就好了。你可以先试试把batch_size设成1跑几十个step看显存曲线,如果还涨就基本排除数据侧问题,重点查模型内部有没有在forward里保存了中间变量或者用了no_grad后又重新requires_grad的操作。另外torch.cuda.memory_summary()确实能看缓存分配,但更推荐用nvidia-smi配合py-spy每隔几秒采样一次,定位涨显存的时间点跟哪个张量相关。
用pdb在每步后打点看内存增量,多半是计算图没释放,检查下有没有把loss或输出存进list里。
这问题我遇到过,大概率不是梯度累积,而是backbone里某些操作在反向传播时把中间变量全保留下来了,比如ResNet的残差连接和ASPP里的空洞卷积。你可以试试在forward里用torch.cuda.memory_allocated()打点,看每个block前后显存涨多少,比memory_summary直观。另外,如果用了SyncBN或者自定义loss里有额外计算,也会导致图不释放,建议把optimizer.zero_grad()放到forward之前试试。我之前是发现模型里一个辅助loss分支忘了detach,导致loss回传路径上有变量被反复保存,去掉就好了。
这问题我碰到过,你试试把dataloader的num_workers设成0跑一下,有时候数据加载线程会偷偷占显存。另外用torch.cuda.memory_summary()抓两次快照对比,重点看是不是有tensor在反向传播后没释放,比如网络里某个分支的输出被保存下来当feature用了。还有个小技巧,在每次迭代后打印一下torch.cuda.memory_allocated(0),如果它一直单调上升而不是波动,那基本就是代码里有变量在累积,跟batch_size关系不大。我之前还遇到过ResNet的BN层在训练模式下统计量更新导致的假性增长,但那个一般涨到一定程度就平了,不会一路冲到OOM。
遇到过一模一样的坑,最后发现不是模型的问题,是dataloader里有个数据增强的操作把计算图给带上了。你试过用torch.cuda.memory_summary()看allocated和reserved的差值吗?如果reserved一直涨而allocated稳定,那大概率是碎片化或者缓存问题,可以试试torch.cuda.empty_cache()放每个epoch开头看看有没有改善。另外DeepLabV3+的ASPP模块里如果用了多尺度特征,有些实现会隐式地保存中间变量用于反向传播,显存自然就线性涨了,建议把requires_grad=False的tensor都显式detach一下。还有一个冷门但常见的原因:验证集或者测试阶段你忘了包torch.no_grad(),虽然batch小但每轮都在积累计算图,显存就会慢慢爬。工具方面除了memory_summary,可以用pytorch的torch.profiler看每层内存分配,或者直接跑一个固定步数然后看显存曲线是阶梯状还是平滑上升——平滑上升多半是缓存没释放,阶梯状就是有东西在累积。你提到没开梯度累积,但确认过optimizer.zero_grad()是每个step都调用了吗?有时候多卡或者自定义训练循环里漏了这一行,反向传播的梯度会一直挂在图上。最后建议跑个最小复现:把backbone换成ResNet18,输入再缩到128,如果显存还涨,那基本就是训练逻辑的问题,跟模型结构无关。
试试关掉amp混合精度,开启gradient_checkpointing,大概率能压住涨势。显存一直涨多半是计算图没释放,用torch.cuda.memory_summary抓一下最稳。
这种情况大概率不是变量泄漏,而是PyTorch的显存缓存机制在作祟,缓存块不会自动释放,看起来就像一直涨。你可以试试在每次迭代后调torch.cuda.empty_cache()看显存会不会回落,如果回落了说明是缓存碎片问题,不影响训练但能缓解OOM。另外用nvidia-smi看显存曲线不如直接在代码里盯torch.cuda.max_memory_allocated(),那个能告诉你峰值到底花在哪。我之前遇到过类似情况,结果是数据加载时num_workers开太多,每个worker都复制了一份模型参数,降成2就正常了,你也可以检查下这块。
查一下是不是backbone的BN层在训练模式下的缓存增长,用torch.cuda.memory_summary看分段最准。
试试把验证集的forward也包在torch.no_grad里,我之前就是这漏了导致显存一路涨。
你这情况先查下dataloader的num_workers,有时候数据加载线程也会占显存,用memory_summary看看是不是缓存没清。
遇到过类似情况,最后发现是dataloader的num_workers开太多导致数据加载时copy到GPU的临时缓存没释放,试试把num_workers调成0或者pin_memory关掉看下。另外你那个每步都涨的曲线,可以打印一下每步的torch.cuda.max_memory_allocated()对比实际分配,如果差值越来越大多半是缓存碎片或者有个隐藏的list在累积tensor。还有个小技巧,用torch.autograd.detect_anomaly()能定位到具体哪一行产生异常梯度,比mem_summary直观。最后检查下是不是验证集也走了forward但忘了with torch.no_grad(),这个最容易忽略。
我之前也踩过类似的坑,你试试在每次optimizer.zero_grad()之后加一行torch.cuda.empty_cache(),有时候是PyTorch的缓存机制没释放,看着像泄露实际是碎片化。另外重点检查一下你的dataloader有没有开num_workers,如果数据加载线程里也碰了CUDA张量,会莫名吃掉显存。还有个土办法,把模型输入固定成torch.randn,然后逐层跑一遍,用torch.cuda.memory_allocated()打印峰值,能快速定位是哪一层爆的。如果确认没循环没detach问题,也可以看看是不是用了BN的track_running_stats=False,这玩意儿在训练时会存一堆中间统计量。
我之前也踩过类似的坑,显存不降反涨基本不是batch_size的问题,更像是某个操作在计算图里被保留了。你试试把torch.no_grad()包住验证集的forward,或者检查一下是不是在循环里反复调用了model(x)却没有把loss清零,尤其要留意有没有把中间变量append到list里。torch.cuda.memory_summary()确实能看allocated和reserved的差值,但更直接的办法是开一个profiler,torch.profiler可以按操作看显存峰值,我上次就是靠它发现是upsample层在反复申请临时buffer。另外,DeepLabV3+的ASPP里有并行分支,如果用了torch.cat,记得确认有没有把原始feature map提前del掉,不然梯度链会一直拽着中间张量不放。还有个偏门但常见的原因,就是dataloader的num_workers设太高了,每个worker预加载的tensor会暂存在CPU端,虽然不占GPU显存,但会间接拖慢backward释放的速度。你也可以试试在每次iteration结束后手动调一下torch.cuda.empty_cache(),虽然治标不治本,但能帮你快速确认是不是有周期性泄漏。如果还找不到,就把loss.item()改成loss.item()之前先断开图,比如loss = loss.detach()再打印,看看显存曲线会不会变平。最后提醒一下,ResNet101的BN层如果没设momentum,training和eval切换时也可能累积状态导致显存异常,但这个比较少见。
我之前也踩过这个坑,最后发现是验证集里忘了包torch.no_grad(),每轮eval都在累积计算图,显存当然只涨不降。你先把训练和验证的forward都检查一遍,确保所有推断路径都关了梯度。另外可以用torch.cuda.memory_summary()看是哪个tensor占着内存,配合pytorch的memory_profiler插件能定位到具体行。如果模型里没有循环,那大概率不是forward多次的问题,倒是可以留意下是不是有变量被存进了list或者dict里没释放。
我之前也踩过类似的坑,DeepLabV3+这种带ASPP的多尺度结构本身就容易在特征拼接时产生大量的中间变量,而且Res101的backbone在低batch下前向计算图占用的显存其实也不小。你先别急着怀疑梯度累积,试着把torch.no_grad()包住验证集的forward看看,如果显存不再涨那基本就是训练时反向传播保存的中间激活值在累积。另一个很隐蔽的点是如果用了类似滑动平均或者EMA的模块,有些实现会额外保存一份模型参数的副本,而且每轮更新时如果没做detach,计算图会越挂越长,显存当然会线性增长。推荐你用pytorch的torch.cuda.memory._record_memory_history()结合snapshot可视化工具,能直接看到每个tensor的分配栈,比memory_summary直观多了。另外检查一下dataloader的num_workers,如果开了多个worker并且每个worker都持有独立CUDA上下文,有时候也会造成显存虚高,但那种情况一般是启动时就高而不是缓慢上涨。还有一个骚操作是直接在损失函数后面加一句loss.backward(retain_graph=False)然后每次优化器step后手动把optimizer.zero_grad()放前面,排除是不是某个自定义loss里不小心用了多次backward。如果这些都没问题,那就看看是否有在循环里反复调用model(x)但没把结果赋给新变量,比如列表推导式里收集了多个输出,每个输出都保留着计算图引用。
我之前也踩过类似的坑,DeepLabV3+这种带ASPP的多分支结构特别容易出问题。你提到显存从2G缓涨到12G,这其实不太像纯显存泄漏,更像是某个tensor在计算图里被意外保留了。建议先跑一个固定step的循环,每个step后打印torch.cuda.memory_allocated(),看是不是线性增长——如果是,那大概率是计算图没释放。我遇到过的情况是,在验证集上忘了关torch.no_grad(),导致每个batch都累积了梯度图,显存就一路涨。另外,ResNet101的BN层在batch_size=2时统计量不稳定,但这不会导致显存涨,所以先排除这个。你确认下是不是用了类似“output, _ = model(x)”这种写法,有时候返回的aux_loss或者中间特征被外部引用也会卡住显存。还有个实用技巧是分段跑,比如forward一次,backward一次,然后手动del loss和output,再torch.cuda.empty_cache(),看显存是否能回落。如果还不行,就用torch.autograd.detect_anomaly(),它会告诉你具体是哪个op在反向传播时产生了异常梯度,很多隐性泄漏都是这么查出来的。至于可视化每层显存,pytorch的memory_summary确实够用了,但更推荐用nvidia-smi的PID监控配合py-spy dump一下python堆栈,能直接看到是哪个线程在持有GPU内存。
这种情况大概率不是真泄露,而是计算图中的历史变量被保留了。你试试在每次反向传播后手动调用optimizer.zero_grad(),同时检查一下有没有把中间特征图存进列表或dict里没释放。另外用torch.cuda.memory_summary()看allocated和reserved的差值,如果reserved特别大说明缓存碎片化,可以试试torch.cuda.empty_cache()。我之前遇到过类似问题,最后发现是模型里一个类似辅助loss的模块在每次forward时创建了新的计算图,建议你把所有可能的可疑节点都加个del再手动gc.collect()看看变化。
torch.cuda.memory_summary()确实得先跑一下,重点看是不是有大量的“allocated”但没被释放的缓存块。我之前遇到过类似情况,最后发现是DataLoader的num_workers开太多,每个worker都在复制模型参数,显存直接翻倍,你这个batch_size都2了还涨,不如试试把worker设成0跑一次。另外DeepLabV3+的ASPP模块里如果用了空洞卷积,某些实现会隐式创建计算图,建议检查一下有没有把中间特征不小心存进了list或者dict里,尤其是循环里append了tensor,那个很坑。
这问题我踩过类似的坑,如果loss和batch都没问题,大概率是有个变量在循环里被反复拼接或者保存了,比如在验证集上跑的时候没有关梯度,或者某个tensor被存进list里了。你试试用torch.cuda.memory_summary()看是哪个阶段涨的,再配合torch.no_grad()包住验证逻辑,如果还涨就检查下是不是有类似loss.item()被赋值给全局变量的操作。另外DeepLabV3+的ASPP模块本身也吃显存,但不会线型涨,所以大概率不是模型结构问题。
我之前也踩过类似的坑,说个最容易被忽略的点:你确认过dataloader的num_workers和pin_memory吗?有时候数据加载线程的预取缓存会占显存,尤其图像分割这种多类别mask的加载,内存拷贝会偷偷累积。另外torch.cuda.memory_summary()确实该跑一下,重点看allocated和reserved的差值,如果reserved一直涨但allocated稳定,大概率是缓存碎片问题,不是真泄漏。还有个土办法,在训练循环里每隔几步打印一次max_memory_allocated,配合nvidia-smi的波动看,能分辨是单步峰值上涨还是跨step累积。至于多次forward,DeepLabV3+本身没有循环结构,但如果你在验证阶段也开了grad,或者用了自定义的aux_loss分支没有及时释放中间tensor,确实会涨。我上次就是忘了在验证时包with torch.no_grad(),结果验证集显存越跑越高。建议你先把batch_size调成1跑100个step,如果显存还是线性涨,就逐个模块forward试,或者用torch.profiler看每层分配,这比纯靠猜靠谱得多。