猴哥的第 51 期分享，欢迎追看
前几天分享的都是 AI 工具系列的文章，今天给大家分享一篇实战方面的笔记。
最近在用 Pytorch 训模型的过程中，发现总是训练几轮后，出现显存爆炸 out-of-memory 的问题，询问了 ChatGPT、查找了各种文档。。。
在此记录这次 debug 之旅，希望对有类似问题的小伙伴有一点点帮助。
训练过程中，网络结构做了一些调整，forward 函数增加了部分计算过程，突然发现 16G 显存不够用了。
用 nvidia-smi 观察显存变化，发现显存一直在有规律地增加，直到 out-of-memory。
尝试思路1：
计算 loss 的过程中是否使用了 item() 取值，比如：
train_loss += loss.item()
我的训练过程不存在这个问题，因为 loss 是最后汇总计算的。
尝试思路2：
训练主程序中添加两行下面的代码，实测发现并没有用。
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
这两行代码是干啥的？
大白话：设置为 True，意味着 cuDNN 会自动寻找最适合当前配置的高效算法，来获得最佳运行效率。这两行通常一起是哦那个
所以：
- 如果网络的输入数据在尺度或类型上变化不大，设置
torch.backends.cudnn.benchmark = True可以增加运行效率； - 如果网络的输入数据在每次迭代都变化，比如多尺度训练，会导致 cnDNN 每次都会去寻找一遍最优配置， 这样反而会降低运行效率 。
尝试思路3：
及时删除临时变量和清空显存的 cache，例如在每轮训练后添加：
torch.cuda.empty_cache()
依旧没有解决显存持续增长的问题，而且如果频繁使用torch.cuda.empty_cache()，会显著增加模型训练时长。
尝试思路4：
排查显存增加的代码位置，既然是增加了部分代码导致的显存增加，那么问题肯定出现在这部分代码中。
为此，可以逐段输出显存占用量，确定问题点在哪。
举个例子：
print("训练前:{}".format(torch.cuda.memory_allocated(0)))
train_epoch(model,data)
print("训练后:{}".format(torch.cuda.memory_allocated(0)))
eval(model,data)
print("评估后:{}".format(torch.cuda.memory_allocated(0)))
最终方案：
最终发现的问题是：我在模型中增加了register_buffer：
self.register_buffer("positives", torch.randn(1, 256))
self.register_buffer("negatives", torch.randn(256, self.num_negatives))
但register_buffer 注册的是非参数的 Tensor，它只是被保存在模型的状态字典中，并不会进行梯度计算啊。
为了验证这一点，还打印出来验证了下：
# for name, param in model.named\_parameters():
for name, param in model.named_buffers():
print(name, param.shape, param.requires_grad)
# 输出如下：
positives torch.Size([1, 256]) False
negatives torch.Size([256, 20480]) False
但是这个buffer却是导致显存不断增加的罪魁祸首。
为此，赶紧把和buffer相关的操作放在torch.no_grad() 上下文中，问题解决！
@torch.no_grad()
def dequeue_samples(self, positives, negatives):
if positives.shape[0] > 0:
self.positives = 0.99*self.positives + 0.01*positives.mean(0, keepdim=True)
self.negatives[:, self.ptr:self.ptr+negatives.shape[1]] = F.normalize(negatives, dim=0)
with torch.no_grad():
keys = F.normalize(self.positives.clone().detach(), dim=1).expand(cur_positives.shape[0], -1)
negs = self.negatives.clone().detach()
如果是训练过程中显存不断增加，问题大概率出现在 forward 过程中，可以通过尝试思路4逐步排查出问题点所在，把不需要梯度计算的操作放在torch.no_grad() 上下文中。
如果本文对你有帮助，欢迎点赞收藏 备用！你的支持是我创作的最大动力。
👇 关注猴哥，快速入门AI工具
# AI 工具：
【玩转云服务】Oracle Cloud甲骨文永久免费云服务器，来薅资本主义羊毛了！白嫖党请进！
Windows上安装Linux子系统，搞台虚拟机玩玩，手把手教
# AI应用** ：**
手把手搭建微信机器人，帮你雇一个24小时在线的个人 AI 助理（上）
视频生成 EasyAnimate-v3 实测，阿里开源，5 分钟带你部署体验，支持高分辨率超长视频
# 效率工具：