记一次Pytorch内存泄露的排查与处理

1.起因:

模型训练过程中内存占用不断增加,训练到30000轮左右已经占用到200G内存.

2.排查过程

查找了网上的一些内存泄漏排查方法,使用了memory_profiler objgraph  pympler这三个工具进行排查

参考链接如下:

Pytorch超出内存

pytorch内存泄漏分析案例 | list转tensor

memory_profiler的使用

python 内存监控模块之memory_profiler

Python 使用objgraph 定位内存泄露

排查python内存泄露中几个工具的使用

记一次python内存泄露的解决过程

pympler跟踪内存泄漏

最终发现内存泄露可能的原因是通过pympler工具,在tker.print_diff()时发现网络训练循环结束后,自定义的两个loss层的自动回传对象仍然未被释放.

3.解决

最终发现问题出在train_net.py中,loss += loss_iter语句,由于loss_iter是一个具有autograd历史的可微变量,因此loss直接加上loss_iter会导致autograd历史不断被保存,内存未被释放.

修改为loss += loss_iter.item()仅仅使用loss_iter数值后问题解决,运行程序在终端使用top命令与free -h命令发现内存不再增加了.