pytorch中碰到的memory leak问题

最近碰到pytorch分布式训练时候，memory几乎线性增加，撑炸机器的问题。
pytorch中内存泄漏常见的原因大概是以下几点：

1. 不恰当的loss累加

• 有些人累加梯度会直接把梯度拿过来加，但是由于每个梯度都存储了很多东西，导致随着step增加，梯度累计的越来越多，造成了内存泄漏
• 做法就是把loss的数值取出来，而不是累计整个梯度

2. 直接把list转化成tensor

• 常出现在dataloader的dataset类实现上，data是list，在get_item的时候直接转化成tensor了，这样好像每次都会造成数据多覆盖。不过这个好像是python天然的原因，算不上pytorch的锅
• 标准做法就是dataset类中存储的是np类型，然后再转tensor(list->np->tensor)

3. dataloader中num_worker的设置

• 有些版本中，num_worker设置大于0，就会有内存泄漏，改成0就没问题了
  参考链接：https://github.com/pytorch/pytorch/issues/13246

有时候可能会遇到不同的问题，具体问题可以通过python的内存分析工具做分析（不过讲道理不是太管用）比如：https://www.pythonf.cn/read/108519，https://zhuanlan.zhihu.com/p/121003986

我的心情随着第一个github的issue答案起起伏伏，试了几遍都不行，然后忽然想到，这些bug官方都回复修了，怎么还能有问题呢…然后转头把sagemaker上pytorch的版本从1.6降到了1.5，世界安静了…

image.png

最近一天一个bug，踩坑美滋滋