【PyTorch笔记】训练时显存一直增加到 out-of-memory？真相了！

最近用 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() 上下文中。

如果本文对你有帮助，欢迎点赞收藏备用！