使用数据加载器时如何从内存中加速批量大小的数据

问题描述

我正在尝试使用 DataLoader 进行训练。数据集150G，都是.npz文件。由于内存大小的限制，一次只能从磁盘读取一个样本。以下是部分代码。

class VimeoDataset(Dataset):
def __init__(self,mode,batch_size=32,num_workers = 8,num_gpus = 4):
    self.batch_size = batch_size
    self.num_workers = num_workers
    self.num_gpus = num_gpus
    self.mode = mode
    self.load_data()
    self.h = 256
    self.w = 448
    xx = np.arange(0,self.w).reshape(1,-1).repeat(self.h,0)
    yy = np.arange(0,self.h).reshape(-1,1).repeat(self.w,1)
    self.grid = np.stack((xx,yy),2).copy()
    self.npzs=[]

    count = self.batch_size * self.num_workers * self.num_gpus
    if self.mode == 'train':
        filelist = glob('/data/vimeoFlow2/dataset/train/*.npz')
        self.npzs = [filelist[i:i + count] for i in range(0,len(filelist),count)]
    else:
        filelist = glob('/data/vimeoFlow2/dataset/val/*.npz')
        self.npzs = [filelist[i:i + count] for i in range(0,count)]

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

def load_data(self,index):
    self.data = []
    self.flow_data = []

    for i in range(len(self.npzs[index])):
        f = np.load(self.npzs[index][i])
        self.data.append(f['i0i1gt'])
        if self.mode == 'train':
            self.flow_data.append(f['ft0ft1'])
        else:
            self.flow_data.append(np.zeros((256,448,4)))    

def getimg(self,index):
    data = self.Meta_data[index]
    img0 = data[0:3].transpose(1,2,0)
    img1 = data[3:6].transpose(1,0)
    gt = data[6:9].transpose(1,0)
    flow_gt = (self.flow_data[index]).transpose(1,0)
    return img0,gt,img1,flow_gt
        
def __getitem__(self,index):        
    img0,flow_gt = self.getimg(index)

dataset = VimeoDataset(mode = 'train',num_gpus = 4)
sampler = distributedSampler(dataset)
train_data = DataLoader(dataset,batch_size=args.batch_size,pin_memory=True,num_workers=args.num_workers,drop_last=True,sampler=sampler)
dataset_val = VimeoDataset(mode = 'val',num_gpus = 4)
val_data = DataLoader(dataset_val,num_workers=args.num_workers)

但是，从磁盘一个一个地读取数据导致数据加载器非常耗时。所以我想改进这个程序，先把num_gpus×num_workers×batch_size的数据量加载到内存中，然后用__getitem__从内存中读取数据，最后每次迭代后替换内存中的数据。但我仍然不知道如何实现它。我在上面的代码中尝试了我的想法。我不知道如何分配 load_data 函数参数。

解决方法

看起来您试图以错误的方式使用火炬 Dataset。 您的 Dataset 子类不应自己批量处理数据，也不应使用工作人员的数量。

批处理数据并并行加载它是 DataLoader 类的作用。您的 Dataset 子类 __getitem__ 方法应该只从数据集中返回 1 个样本（以及另外一个基本事实注释），它应该是像 Tensor 或 Array 这样可以连接的数据以创建批处理。

查看 Dataset 和 DataLoader 文档，其中对此非常清楚。

DataLoader 的目的是并行加载（即从磁盘读取到内存）和预处理数据。如果指定了 8 个 worker，则大致意味着 8 个并行线程正在调用 __getitem__ 方法来创建一批项目。请注意，DataLoader 已经“缓存”了数据并提前加载它们以便及时准备好（查看 prefetch_factor 参数）。

这应该是加载速度和内存消耗之间的充分折衷，您应该在编写任何自定义缓存、加载和并行处理数据之前尝试这样做。

deep-learning python pytorch