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. 首先将模型从流程中剔除,只保留加载数据的部分,看内存是否增长。结果发现内存还是正常增长,确认是数据加载的过程中出现问题。
@profile(precision=4,stream=open('memory_profiler.log','w+'))
def test(dataset, model, device):
    for i, data in tqdm(enumerate(dataset)):
        # -----------------------------------
        # imgs: [batch, channel, frames, H, W]
        # audio: [batch, 1, T, F]
        # -----------------------------------
        imgs, mask, audio = data
        # model.set_input(imgs.to(device), mask.to(device), audio.to(device))
        if i > 200:
            break
  1. 将@profile装饰器放到__getitem__上,查看哪行代码或者哪个函数存在内存增长。发现某些行会出现一定的增长,但不是每次执行都有increment。到这基本上说明该工具派不上用场了。
    (忘了截图了【捂脸】)
Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
   235 3143.3828 MiB 3143.3828 MiB           1       @profile(precision=4,stream=open('memory_profiler.log','w+'))
   236                                             def _audio_transform(self, decord_ar, frame_id_list, fps):
   237 3143.3828 MiB   0.0000 MiB           1           sample_rate = 16000
   238 3143.3828 MiB   0.0000 MiB           1           hop_length = 160
   239 3143.3828 MiB   0.0000 MiB           1           num = sample_rate // fps # 每帧对应的音频特征数
   240 3143.3828 MiB   0.0000 MiB           1           audio_feature_id_list = np.multiply(frame_id_list, num) # 视频帧对应的音频特征
   241                                         
   242 3154.0469 MiB   0.5195 MiB           2           audio_transform = transforms.Compose([
   243 3153.5273 MiB  10.1445 MiB           2               MelSpectrogram(
   244 3143.3828 MiB   0.0000 MiB           1                   sample_rate=sample_rate,
   245 3143.3828 MiB   0.0000 MiB           1                   hop_length=hop_length,
   246 3143.3828 MiB   0.0000 MiB           1                   n_fft=512,
   247 3143.3828 MiB   0.0000 MiB           1                   n_mels=128
   248                                                     ),
   249                                                     # Lambda(
   250                                                     #     lambda x: x[:, :-(x.size(1) - args.dataset.fps * args.audio2video)]
   251                                                     #     if x.size(1) > args.dataset.fps * args.audio2video else x
   252                                                     # ),  # 截断长度为1秒
   253 3161.1797 MiB   7.1328 MiB           4               transforms.Lambda(lambda x: amplitude_to_DB(
   254 3159.9375 MiB   0.0000 MiB           1                   x, multiplier=10, amin=1e-10, db_multiplier=math.log10(max(1e-10, torch.max(x).item())), top_db=80
   255                                                     )),
   256 3161.1797 MiB   0.0000 MiB           3               transforms.Lambda(lambda x: (x + 40) / 40),
   257                                                     # transforms.Lambda(lambda x: x.transpose(1, 0).unsqueeze(0)),  # (F, T) -> (1, T, F)
   258 3161.1797 MiB   0.0000 MiB           3               transforms.Lambda(lambda x: x.transpose(1, 2)),  # (1, F, T) -> (1, T, F)
   259                                         
   260                                                 ])
   261 3154.0469 MiB   0.0000 MiB           1           audio = decord_ar._array
   262                                                 # audio sample according to video
   263 3154.0469 MiB   0.0000 MiB           1           audio_sample_list = np.array([])
   264 3154.0469 MiB   0.0000 MiB           1           if self.new_step == 1: # no skip
   265 3154.0469 MiB   0.0000 MiB           1               audio_sample_list = audio[:,audio_feature_id_list[0]: audio_feature_id_list[-1] + num]
   266                                                 else:
   267                                                     for i in audio_feature_id_list:
   268                                                         audio_sample_list.append(audio[i : i + num])
   269                                                     
   270                                         
   271                                                 # 补齐为指定frames对应的音频长度
   272 3154.0469 MiB   0.0000 MiB           1           if audio_sample_list.shape[1] < self.new_length * num:
   273                                                     audio_sample_list = torch.nn.functional.pad(audio_sample_list, (0, self.new_length * num - audio_sample_list.shape[1]))
   274                                         
   275                                         
   276 3154.0469 MiB   0.0000 MiB           1           audio_sample_list = torch.tensor(audio_sample_list, dtype=torch.float32)
   277 3161.1797 MiB   0.0000 MiB           1           return audio_transform(audio_sample_list)

在音频转换阶段,在后期执行过程的memory_profiler工具没有检查到Increment有增长,说明这部分应该没有内存泄漏。(但前几次执行的执行的时候有出现10.1445MB的增长我也不知道为什么)

step3:于是一步步将部分函数或代码注释,从而缩小定位范围
将代码tranform转换的所有步骤的注释,只保留decord读取音视频部分

    def __getitem__(self, index):
        directory, target, duration = copy.deepcopy(self.clips[index])
        directory = os.path.join(self.video_dir, directory)
        if self.video_loader:
            if '.' in directory.split('/')[-1]:
                # data in the "setting" file already have extension, e.g., demo.mp4
                video_name = directory
            else:
                # data in the "setting" file do not have extension, e.g., demo
                # So we need to provide extension (i.e., .mp4) to complete the file name.
                video_name = '{}.{}'.format(directory, self.video_ext)

            decord_vr = decord.VideoReader(video_name, num_threads=3)
            fps = int(decord_vr.get_avg_fps())
            decord_ar = decord.AudioReader(video_name)
        return (decord_vr, decord_ar, None)

结果发现还是存在内存增长,但是这次的增长规律有所不同。由于没有了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

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

Future:如果上述步骤还不行,还打算将训练集划分为多个子训练集,企图在跳出循环时能释放该dataloader对象的内存;或者是将数据进行预处理存在本地,用硬盘空间换内存空间。

结论

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

评论 3
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值