问题
在训练模型时,使用pip安装的decord库读取视频和音频,但在运行过程中遇到cpu内存泄漏的问题,加载了大约60w个视频样本后就会占用接近300G的cpu内存
解决方案
step1:参考常规的内存泄漏的检查思路,排查代码中可能存在的问题
网上通常出现内存泄漏的解决方法有以下几种:
- loss累积的时候没有使用item(),导致梯度没有正常释放。因此需要在loss积累的地方加上item()
#loss_all += loss
loss_all += loss.item()
#a = torch.tensor(b)
a = torch.tensor(b, dtype=torch.float32)
# data = self.data_list[index]
data = copy.deepcopy(self.data_list[index])
- 使用del和gc.collect()显式清空内存
del data
gc.collect()
step2:以上方法都尝试过,但都没有改变我代码中内存泄漏的问题。因此想着使用memory_profiler工具定位代码中哪个地方导致内存增加(参考)
- 首先将模型从流程中剔除,只保留加载数据的部分,看内存是否增长。结果发现内存还是正常增长,确认是数据加载的过程中出现问题。
@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
- 将@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。