python的decord库存在内存泄漏

问题

在训练模型时,使用pip安装的decord库读取视频和音频,但在运行过程中遇到cpu内存泄漏的问题,加载了大约60w个视频样本后就会占用接近300G的cpu内存

解决方案

step1:参考常规的内存泄漏的检查思路,排查代码中可能存在的问题
网上通常出现内存泄漏的解决方法有以下几种:

  1. loss累积的时候没有使用item(),导致梯度没有正常释放。因此需要在loss积累的地方加上item()
#loss_all += loss
loss_all += loss.item()
  1. list转换成tensor时,指定dtype(参考),以及尽量不适用list,而使用numpy.array(参考1)(参考2)
#a = torch.tensor(b)
a = torch.tensor(b, dtype=torch.float32)
  1. 将num_workers设置为0(参考)
    这个解决方案最直接就是会导致读取速率明显变慢
  2. 对可能被修改的变量进行复制时使用copy()(参考)
# data = self.data_list[index]
data = copy.deepcopy(self.data_list[index])
  1. 使用del和gc.collect()显式清空内存
del data
gc.collect()

step2:以上方法都尝试过,但都没有改变我代码中内存泄漏的问题。因此想着使用memory_profiler工具定位代码中哪个地方导致内存增加(参考)
【图】

  1. 首先将模型从流程中剔除,只保留加载数据的部分,看内存是否增长。结果发现内存还是正常增长,确认是数据加载的过程中出现问题。
  2. 将@profile装饰器放到__getitem__上,查看哪行代码或者哪个函数存在内存增长。发现某些行会出现一定的增长,但不是每次执行都有increment。到这基本上说明该工具派不上用场了。

step3:于是一步步将部分函数或代码注释,从而缩小定位范围
将代码tranform转换的所有步骤的注释,只保留decord读取音视频部分
【代码】
结果发现还是存在内存增长,但是这次的增长规律有所不同。由于没有了transform步骤,内存增长趋势没有大的波动,没有明显内存释放的过程,而是缓慢地一步步的增长。到这里基本确定是decord库可能存在内存泄漏问题。

step4:去github上搜decord代码中存在的memory leak问题的Issue。参考1、参考2
从一个回答中找到了问题所在: The problem appears with ctx=cpu(0) but not ctx=gpu(0)
因此我试了以下使用ctx=gpu(0)来读取,结果发现我的decord不支持decord.gpu(0)。然后发现别人的decord都是从源代码编译安装的,而我的是从pip安装来的。于是,尝试去https://github.com/dmlc/decord用源码安装decord

# 下载编译
git clone --recursive https://github.com/dmlc/decord
cd decord
mkdir build && cd build
cmake .. -DUSE_CUDA=ON -DCMAKE_CUDA_COMPILER=/path/to/cuda/bin/nvcc -DCMAKE_BUILD_TYPE=Release
make

# 安装python
cd ../python
# option 1: add python path to $PYTHONPATH, you will need to install numpy separately
pwd=$PWD
echo "PYTHONPATH=$PYTHONPATH:$pwd" >> ~/.zshrc
source ~/.zshrc
# option 2: install with setuptools
python setup.py install --user

结果发现内存不动了!!至此问题解决~~!!

结论

使用pip安装的decord库读取音视频时,存在内存泄漏的问题,并且这个版本的decord缺少部分功能。解决方法是使用decord源码重新编译安装到python。

你可能感兴趣的:(内存泄漏,decord,python)