用PyTorch训练模型时，显存突然暴涨，怎么排查具体是哪一行代码导致的？

试试torch.cuda.memory_allocated()在可疑的transform前后打点，或者用line_profiler配torch.cuda看每行增量，比summary直观多了。

显存跑到第5个epoch才爆，基本能排除模型结构本身的问题，更可能是数据增强里某些操作把计算图带进来了。你可以先查查自定义Dataset的__getitem__里有没有在transform中调用需要梯度的操作，比如不小心对tensor做了in-place修改或者保留了中间变量引用。torch.cuda.memory_summary()确实不够直观，我更常用torch.cuda.memory_allocated()配合手动打点，在每个epoch开始和结束各记一次，看增量出现在哪个阶段。另外可以试试torch.cuda.reset_peak_memory_stats()在每个epoch前重置，然后打印峰值，这样能快速锁定是哪个epoch开始异常。如果怀疑是某一行，用pdb在可疑位置断点，配合nvidia-smi实时看显存变化也挺有效。还有个偏方是把DataLoader的num_workers设成0，排除多进程导致的显存统计混乱。实在不行就用torch.autograd.set_detect_anomaly(True)跑一小段，虽然慢但能揪出反向传播里的异常节点。

试试torch.cuda.memory_allocated配合逐行断点，八成是某个增强里张量没detach一直在累积。

这个问题我还真踩过类似的坑，感觉大概率不是传统意义上的内存泄漏，而是某个数据增强把中间张量留在了计算图里，导致反向传播时显存越积越多。你可以先试试torch.cuda.memory_allocated()在每个epoch前后打印一下，看是持续增长还是某个batch突然跳上去的。如果是一跳一跳的，八成是transform里用了torch的张量操作但没加with torch.no_grad()，尤其是你新加的那几个增强。另外我强烈推荐装个torch.cuda.memory._record_memory_history()，然后在爆显存的时候dump出快照，用pytorch官方的memory visualizer网页打开，能直接看到是哪一行分配了最多显存，比memory_summary直观多了。还有个小技巧，把DataLoader的num_workers设成0单进程跑一遍，如果显存正常了，那就是worker里某个增强或者dataset缓存没清。cProfile只能看时间，看显存得靠上面那个history recorder或者第三方包如pytorch_memlab的line_profiler。你要实在懒，就在可疑的transform前后插torch.cuda.empty_cache()加打印，虽然土但定位快。

你可以试试 torch.cuda.memory_snapshot() 配合 pickle 把快照存下来，然后用 PyTorch 官方那个 memory visualizer 网页工具看，能直接定位到是哪次分配炸的。我上次也是加了个随机旋转的 transform，结果每个 epoch 都偷偷缓存了一堆 tensor 没释放。另外记得检查下自定义 Dataset 里有没有把整个数据集读进内存当成员变量，那种坑特别隐蔽。实在不行就二分法注释 transform，跑两个 epoch 看显存曲线，比啥工具都直接。

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