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

我之前也遇到过类似情况，加了个随机裁剪之后显存直接翻倍，最后发现是transform里对同一张图重复做了多次GPU张量运算，没及时转回CPU。建议你先把数据增强里所有涉及to(device)的操作全部去掉，强制用CPU张量做预处理，如果显存正常了那就是这块的问题。另外你说的按行看显存，torch没有直接等价cProfile的工具，但可以用torch.profiler配合memory profiling，能输出每个张量分配时对应的python调用栈，虽然有点重但够精准。还有个土办法，就是在每个epoch里手动打点，用torch.cuda.max_memory_allocated()对比前后差值，用二分法注释掉部分代码，定位到具体模块很快。另外注意一下Dataset的__getitem__里如果有累积list或者没释放的中间变量，也会造成缓慢泄漏，但你这个是突然暴涨，更像某个batch里异常大的输入导致，可以检查下是不是数据增强产生了极端尺寸的图。最后推荐个第三方库叫pytorch_memlab，它有个line-by-line的decorator，能直接标注函数内每行的显存增量，比memory_summary直观多了。

之前也遇到过类似情况，后来发现是数据增强里用了随机crop没限制边界，导致每次迭代张量形状不一样，PyTorch底层缓存没清干净。可以先试试把transform里的操作逐个注释掉跑一个epoch，用nvidia-smi实时盯显存，哪个操作一取消显存就稳了基本就是它。另外torch.cuda.empty_cache()放在每个epoch结尾清一下缓存，虽然不解决根本问题但能缓解，配合memory_allocated和max_memory_allocated对比看差值能缩小范围。真要按行定位，可以试下torch.profiler或者pytorch_memlab，虽然没cProfile那么直观，但能打印每个张量的分配位置，比memory_summary好懂些。

试试torch.cuda.memory._record_memory_history()配合tracemalloc，能按行看分配栈，比summary直观多了。

我之前也遇到过这种加了transform之后显存爆炸的情况，最后发现是某个随机裁剪操作里生成了大量中间变量没及时del，你试试用torch.autograd.detect_anomaly()或者干脆把每个transform单独跑一遍对比显存峰值，定位会快很多。另外可以试试pytorch的torch.profiler，它支持按行显示显存分配，虽然配置起来有点麻烦但比memory_summary直观多了。还有个小技巧，把数据加载的num_workers设成0跑一两个batch，能排除掉DataLoader里内存泄漏的干扰。

我前几天也踩过类似的坑，加了个随机裁剪的transform，结果显存从8G直接飙到20G，后来发现是裁剪边界处理时生成了超大的临时张量，建议你优先检查transform里有没有用到类似torch.where或者mask索引的操作，这些经常会隐式创建中间变量。另外那个memory_summary确实不好用，我现在的习惯是直接在可疑代码块前后打两行print(torch.cuda.memory_allocated())，差值一算就能锁定范围，虽然土但效率很高。如果你愿意装第三方库的话，可以试试pytorch_memlab，它的line_profile装饰器能按函数行号输出显存占用，比cProfile直观很多。还有个小细节，如果你在Dataset里用了albumentations，记得检查它的返回值是不是numpy数组，有时候忘了转tensor会反复拷贝到GPU显存里。最后提醒下，数据增强最好是放在CPU上做，别一股脑全塞进GPU，不然预处理峰值很容易超过训练峰值。

试试pytorch的torch.autograd.detect_anomaly，能定位到loss反传时爆显存的具体op，数据增强的话查查transform里有没有缓存没用del。

可以把batch size调成1跑一遍，用nvidia-smi dmon实时盯显存，哪个transform前后涨了就是哪的问题。

这种问题我也踩过坑，加了几个transform之后显存直接翻倍，最后发现是随机裁剪那块在每次迭代里都重新创建了索引数组，没释放。你可以试试用pytorch的autograd检测工具，或者干脆开一个子进程跑训练，用nvidia-smi实时监控显存变化，再配合二分法注释代码段，比看memory_summary直观多了。另外检查下是不是在Dataset里用了list存储中间结果忘了清，我之前就是被这个坑的。

试试torch.cuda.set_per_process_memory_fraction配合二分法，或者用nvidia-smi盯显存的同时加print逐段缩范围，爆点通常在新加的transform里。

给代码每行标号，用pdb逐步跑，显存涨的时候看是哪步分配的张量没释放，比盲猜快。

我之前也遇到过类似情况，后来发现是DataLoader的num_workers开太多，加上transform里用了随机裁剪导致张量不停累积，试试把num_workers调小或者把transform的中间结果转成numpy存一下，显存会稳很多。而且内存泄漏不一定是显存，也可能是CPU内存爆了拖累GPU，建议用tracemalloc看下Python对象分配。真要按行定位的话，可以试试pytorch的torch.autograd.detect_anomaly，但那个是查梯度的，对显存用处不大，还是得靠手动二分注释代码来排除。

我之前也踩过类似的坑，加了几个在线随机crop和rotation之后，显存直接翻倍，最后发现是transform里对tensor调用了.cpu()然后又转回cuda，导致每次迭代都创建了新的中间变量。你那个情况不一定是泄漏，很可能是某个操作把计算图给保留了，比如在Dataset的__getitem__里用了torch.tensor()而不是torch.from_numpy()，这样梯度虽然没开，但autograd还是会记录一些临时节点。定位的话，别死磕memory_summary，试试torch.cuda.set_per_process_memory_fraction把显存限制到刚好够用，然后跑的时候用faulthandler或者pdb在OOM前打断点，看堆栈最后几帧就能锁定。更实用的是用py-spy的dump --locals看每个Python栈的局部变量，配合nvidia-smi的实时显存变化，基本能判断是哪个循环在累积。另外你检查下数据增强里有没有用random.randint这种，如果没用固定种子，每个epoch的随机性会导致某些batch特别大，加上之前没释放的缓存就爆了。还有个土办法，把增强逐步注释掉跑一个epoch，二分法找罪魁祸首，比看工具快。如果怀疑是内存没释放，就在__getitem__末尾手动del掉大tensor，再torch.cuda.empty_cache()试试，但注意这招只能治标，真正问题是某个算子返回了不可变视图导致原buffer没被回收。

试试给每个transform前后插个torch.cuda.synchronize看峰值，八成是随机crop那块没释放中间变量。

用pytorch的memory profiler钩子，能按张量分配点回溯到代码行，比summary好用。

试试用pytorch的autograd记录每个张量的生命周期，或者直接拆掉transform逐个跑一遍，比看summary直观多了。

用torch.cuda.set_per_process_memory_fraction限制显存，再配合tracemalloc查Python对象分配，基本能锁死泄漏点。

试试pytorch_memlab的LineProfiler，能按行看显存，比memory_summary直观多了。

试试用torch.profiler带profile_memory=True按行看，或者跑个batch用tracemalloc查Python侧显存。

显存暴涨多半是反向传播时激活值没释放，先用小batch逐步二分定位到具体模块。

试试用pytorch的autograd记录每个张量的峰值，配合torch.cuda.set_per_process_memory_fraction设个上限，爆了之后回看backward的hook，基本能定位到是哪个op在反向时分配了巨量显存。另外你那几个数据增强是不是用了.numpy()或者.item()？这些操作会把tensor从GPU搬到CPU，如果没及时del，内存会越积越多。还有个笨办法：把transform一个个注释掉跑一个epoch，看哪个加上去显存曲线斜率不对，比看summary直观多了。

试试用pytorch的memory_profiler，能按行看显存，或者先注释掉新加的transform逐个排除。

我上次也这样，结果是一个随机裁剪的缓存没清，加上del和empty_cache就稳了。

试过用torch.cuda.memory._record_memory_history()配合torch.cuda.memory._snapshot()吗？这俩能生成一个完整的分配栈，然后你用chrome://tracing打开导出的json，能直接看到每个tensor的分配调用链，比memory_summary直观多了。不过要注意开启后会拖慢训练速度，建议专门跑一个debug epoch来抓。

另外你提到数据增强，我猜大概率不是显存泄漏，而是某个transform把输入张量变成了不必要的大中间变量。比如随机裁剪如果用了torchvision.transforms.RandomResizedCrop，它的参数scale范围设太宽，可能偶尔生成超大尺寸的feature，反向传播时梯度也跟着爆炸。可以先试试把所有transform改成确定性操作，看还爆不爆。

如果非要用现有代码，有个土办法：在DataLoader的collate_fn里加一行torch.cuda.synchronize()然后打印torch.cuda.memory_allocated()，每取一个batch记录一次，很快就能看出是数据加载时涨还是forward/backward时涨。还有个小技巧，把torch.no_grad()包在数据预处理那段，能强制切断一些隐式计算图的累积。

要是确认是自定义Dataset的问题，检查下有没有在__getitem__里把张量往cuda上搬，或者用全局变量存了中间结果。我之前遇到过类似情况，最后发现是某个opencv操作返回的numpy数组没转成float32，导致后续转tensor时隐式复制了多次。

我之前也遇到过类似的，后来发现是数据增强里某个操作在GPU上执行了，结果每步都留了计算图没释放。你可以试试用torch.autograd.detect_anomaly()开一下，它会提示反向传播时哪步出了问题，虽然不能直接定位到行，但能缩小范围。另外，建议把transform里的操作统一放到CPU上做，或者用torch.no_grad()包住预处理部分，大概率能解决你的问题。

用torch.cuda.memory._record_memory_history()配合snapshot工具，能按行看分配堆栈，比summary直观多了。

我之前也遇到过类似情况，最后发现是数据增强里某个操作在GPU上动态创建了超大临时张量，比如随机缩放时用了F.interpolate没指定align_corners，导致中间变量没被释放。你可以试试torch.autograd.detect_anomaly()，它能直接定位到产生NaN或异常梯度的前向代码行，虽然不一定直接报显存，但往往能顺藤摸瓜找到问题源头。另外建议把Dataset里的每个transform单独拆出来跑一遍，用nvidia-smi盯着看哪个操作峰值暴涨，比看memory_summary直观多了。