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

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 子类既不应该对数据本身进行批处理，也不应该使用工人的数量。

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

查看Dataset 和DataLoader 文档，这方面的内容非常清楚。

---

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

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