PyTorch模型加载到一半显存爆了，是不是我数据预处理有毛病？

验证时不更新BN也会占显存，试试torch.inference_mode()，比no_grad更省。

验证阶段爆显存跟数据预处理关系不大，pin_memory那点开销也不至于。你提到的eval模式确实要检查，BN在train模式下会存batch统计量，验证时还占着计算图就麻烦了。不过更可疑的是你加载了完整checkpoint，optimizer状态里动量缓冲会额外占一份模型大小的显存，验证时根本用不上，建议只load模型权重试试。另外检查下验证循环里有没有不小心把loss或输出累加进list没释放，这个坑挺常见的。

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

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

验证阶段爆显存但训练不爆，这个现象其实挺常见的，你提到的几个方向里我觉得最可疑的还是checkpoint里带了optimizer状态。你加载完整checkpoint时，optimizer的动量buffer和梯度的历史状态都会被恢复，这部分在验证时完全用不上，但确实会额外吃掉不少显存。建议验证时只load model的state_dict，optimizer那部分直接跳过，或者单独存一个纯权重的版本。至于eval()模式，确实要记得切，但BN层在no_grad下即使不切也不会存梯度，主要影响的是统计量的更新行为，对显存影响没那么大。pin_memory的话主要是锁页内存占用，跟显存关系不大，但验证时确实没必要开。还有一个容易被忽略的点是验证时前向的输出有没有及时释放，比如你累积了预测结果做指标计算，如果一直挂在计算图上或者list里存着不放，累积几个batch就上去了。另外可以试试在验证循环里加torch.cuda.empty_cache()看看显存曲线，定位到底是哪一步涨上去的。