PyTorch模型训练时显存爆炸，但batch size已经很小了，还有什么排查思路？

我之前也踩过类似的坑，检查完数据和batch size后才发现是模型里用了太多中间变量，特别是DeepLabV3+的ASPP模块和decoder部分，建议用torch.cuda.memory_stats或者直接hook每层输出看看。另外你的DataLoader里如果对图像做了随机crop或resize，得确认一下实际返回的tensor尺寸是不是真的统一了，有时候多进程worker里缓存了变量也会导致显存峰值。实在不行可以试试torch.utils.checkpoint，把部分层换成梯度检查点，能省不少显存。

我之前也遇到过类似情况，排查下来发现是输入图像没做归一化，直接以uint8喂进模型，导致自动转float时临时张量暴增。你可以试试用torch.cuda.memory_snapshot或nvidia-smi的进程监控对比前后变化，另外检查下forward里有没有重复创建大tensor或者用了不释放的list存中间结果。还有个小技巧，把batch size设成1跑一次，如果显存还是异常高，那基本就是模型结构或数据加载的问题，跟batch关系不大了。

我之前也踩过类似的坑，3090跑DeepLabV3+按理说512分辨率不该这么惨。你先别急着怀疑DataLoader，试试用torch.cuda.reset_peak_memory_stats()和torch.cuda.max_memory_allocated()在每步前后打点，看看峰值到底出现在前向还是反向。我遇到过最隐蔽的问题是模型里用了自定义的注意力或者上采样层，内部创建了临时张量没释放，尤其是那种带循环或者多次调用的子模块，显存会像滚雪球一样涨。另外检查一下有没有不小心把梯度传给了不需要更新的参数，或者用了requires_grad=True的buffer，这会导致反向时保存大量中间激活。工具方面，torch.profiler自带内存分析，能按行号看每行代码的分配量，比手动猜快多了。还有个土办法，把batch size降到1跑通，然后逐步加，同时用nvidia-smi每隔0.1秒记录显存曲线，看是不是线性增长而非突增，如果是线性增长那基本就是某些张量没释放。对了，你试试在dataloader里把num_workers设成0，有时候多进程预取会额外复制CUDA张量，虽然通常不占显存但会影响峰值。真不行就换mixed precision，amp能省将近一半显存，3090对fp16支持很好，代价只是偶尔需要调loss scale。最后提醒下，DeepLabV3+的ASPP模块里空洞卷积的dilation rate大时，feature map的显存占用会超出你预期，可以考虑用深度可分离卷积替换常规卷积来降内存。

我之前也遇到过类似情况，排查下来发现是输入尺寸没对齐，模型里有些下采样层会让特征图尺寸变成非整数，pytorch会自动填充导致显存暴涨，你可以打印一下中间层的shape看看。工具方面推荐用torch.profiler或者pytorch_memlab，能按行显示每层内存占用，直接定位到具体模块。另外检查下dataloader里是不是用了pin_memory和num_workers太高，有时候prefetch机制会额外占不少显存，试试把workers降到0看看还崩不崩。最后如果所有都正常，可以手动用torch.cuda.max_memory_allocated()对比峰值和实际消耗，排除是不是别的程序占着显存没释放。

DataLoader里如果用了num_workers>0还开pin_memory，有时反而会堆积显存，可以先试试设成0排除掉。另外推荐用torch.cuda.memory_summary()看下快照，或者装个torchprof，能按层打印显存峰值。还有个坑是验证阶段忘了加torch.no_grad()，中间变量一直不释放，我踩过好几次。

可以试试 torch.cuda.memory_summary() 看下当前显存分布，再配合 torch.cuda.synchronize() 在训练循环里逐段打点。DeepLabV3+的ASPP模块在512尺寸下中间激活其实挺吃显存的，不一定是DataLoader的锅。另外检查下有没有在forward里保留了不需要的tensor引用，比如list里append了带梯度的输出，这种容易悄悄堆积。

我之前也踩过这个坑，DeepLabV3+在512x512下batch size 4就OOM确实不太正常，3090跑16应该是没问题的。你先别急着怀疑DataLoader，最该看的其实是模型里那些中间变量，尤其是ASPP模块和decoder里的双线性插值，这些地方很容易在backward时把激活值全存下来。可以用torch.cuda.memory_summary()打一下快照，或者装个torchinfo看每层的输出shape和参数量，虽然它不直接给显存，但能帮你找到异常大的中间tensor。另外检查下是不是在forward里做了cat或者多次unsqueeze，这种操作会悄悄复制数据。还有一个容易被忽略的点：如果你用了预训练权重但没freeze backbone，BN层的running stats和梯度也会额外吃显存。建议先跑一个batch，在loss.backward()之后立刻打印显存，看是前向就炸还是反向才炸，这样能快速缩小范围。

你这情况我遇到过，大概率是DataLoader里num_workers开太大加上pin_memory，数据在CPU和GPU之间来回拷贝也会吃显存。另外DeepLabV3+的ASPP模块中间特征图挺占地方的，可以试试torch.cuda.memory_summary()看下峰值分配在哪。还有个容易忽略的点：验证阶段没加torch.no_grad()，或者loss累积时没detach，显存会一直涨。建议先用小尺寸输入跑一个iter，打印每层输出shape和显存，基本就能定位了。

可以试试 torch.cuda.memory_summary() 看下峰值卡在哪个阶段，另外用 torch.cuda.memory_snapshot() 配合 pytorch 的 _record_memory_history 能画出显存时间线，定位是哪一步在涨。我之前遇到过 DataLoader 里 pin_memory 加 num_workers 太大反而把显存吃爆的情况，你可以先把 workers 设成 0 排除一下。还有 DeepLabV3+ 的 ASPP 那几层空洞卷积中间激活挺占地方的，可以试试 torch.utils.checkpoint 把那部分包起来换显存。

先试试 torch.cuda.memory_summary() 看看到底是哪块在涨，十有八九是 DataLoader 里张量没释放。

试试 torch.cuda.memory_summary()，再检查下 DataLoader 的 num_workers 和 pin_memory，有时候是数据增强在 GPU 上偷偷建图了。

DeepLabV3+在512分辨率下batch 4就OOM确实不太正常，大概率不是模型本身的问题。先查DataLoader里有没有默认pin_memory加num_workers开太大，再确认下是不是验证阶段没加torch.no_grad()导致图一直挂着。想看每层占用可以用torch.cuda.memory_summary()或者装个pytorch_memlab，能直接定位到具体哪个模块在堆中间变量。另外把cudnn benchmark关掉试试，有时候算法选择会偷偷多占不少显存。

我之前也踩过这个坑，后来发现是验证阶段没加torch.no_grad()，中间激活全留着，显存直接翻倍。你可以先用torch.cuda.memory_summary()看看峰值卡在哪个阶段，再配合torch.cuda.memory_snapshot()定位具体是DataLoader还是前向传播在吃显存。另外DeepLabV3+的ASPP模块中间特征图挺大的，试试把backbone的output_stride调成16，能省不少。

写得挺好，建议补充一些性能数据。

学到了，感谢分享！

刚入门，这个对我帮助很大。