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

Question

我正在尝试使用数据加载器进行训练。数据集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, len(filelist), 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, 2, 0)
    gt = data[6:9].transpose(1, 2, 0)
    flow_gt = (self.flow_data[index]).transpose(1, 2, 0)
    return img0, gt, img1, flow_gt
        
def __getitem__(self, index):        
    img0, gt, img1, flow_gt = self.getimg(index)

dataset = VimeoDataset(mode = 'train',  batch_size=32, num_workers = 8, 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',  batch_size=32, num_workers = 8, num_gpus = 4)
val_data = DataLoader(dataset_val, batch_size=args.batch_size, pin_memory=True, num_workers=args.num_workers)

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

Answer 1

您似乎试图以错误的方式使用 torch Dataset。 您的 Dataset subclass 既不应该对数据本身进行批处理，也不应该使用工人的数量。

批处理数据并并行加载是 DataLoader class 的作用。您 Dataset subclass __getitem__ 方法应该只从数据集中 returns 1 个样本（以及另外一个基本事实注释），它应该是 Tensor 或Array 可以串联起来创建一个批次。

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

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

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

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

How to accelerate batch-size data from memory when using dataloader

python

deep-learning

pytorch

dataloader