Pytorch DataLoader 当 num_worder>0 时内存占用逐渐增大最终OOM

现象

内存占用随着训练过程逐渐增大,最终Out of memory。即使将不再用到的变量及时删除并使用gc.collect(),也无法解决。

解决方案

方案1:Dataset 的成员变量避免使用Python原生数据格式

以下面的Dataset代码为例,如果使用 self.data = [x for x in range(24000000)],即使用了List作为Dataset的成员变量,当 num_worker>0 时就会导致前面的现象,而改用numpy ndarray 包装一下就没问题了。另外,参考资料3里面说用DataFrame也会有问题。

小结:定义Dataset的成员变量时,使用 numpy ndarray 代替 List 和 DataFrame.

from torch.utils.data import Dataset, DataLoader
import numpy as np
import torch

class DataIter(Dataset):
    def __init__(self):
    	self.data = [x for x in range(24000000)]  # 有内存溢出风险!!!
    	# 改成下面这样可以避免!
        self.data = np.array([x for x in range(24000000)])  
        

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        data = self.data[idx]
        data = np.array([data], dtype=np.int64)
        return torch.tensor(data)

方案2:将 num_worker 设成 < 总CPU物理核心数

# 总核数 = 物理CPU个数 X 每颗物理CPU的核数 

# 查看物理CPU个数
cat /proc/cpuinfo| grep "physical id"| sort| uniq| wc -l

# 查看每个物理CPU中core的个数(即核数)
cat /proc/cpuinfo| grep "cpu cores"| uniq

例如我的机器得到结果是 物理CPU个数 1 乘以 7个物理核/CPU=总7个物理核心,所以需要将 num_worker 设得比7更小,例如可以设为 5(设为5表示有5个子进程,加上主进程,整个训练的过程就启动了6个进程).

方案3:将num_worker设为0

如果前面的方案都不奏效,而且2个一起用上也没用的话,最后还可以考虑将num_worker设为0,设为0表示只使用主进程进行加载数据。

参考资料

[1] https://discuss.pytorch.org/t/num-workers-in-dataloader-will-increase-memory-usage/28522/7
[2] https://github.com/pytorch/pytorch/issues/13246
[3] https://github.com/pytorch/pytorch/issues/5902

你可能感兴趣的:(PyTorch,pytorch,dataloader,内存,内存泄漏)