dataloader读取大量图片如何节省内存空间

在FCN一章时,读取共约2900张图片,当时服务器内存为22GB,现在加上了Pascal VOC2012的增强数据集,总共10000+张图片,服务器会爆内存,想问一下各位大佬,能不能改变图片的读取方式,不一次性把所有数据加载到内存。

母鸡啊,我也是,io太慢了,官方给的代码是先把所有图片读取到内存,服了,每次训练前读半小时的图片

我尝试过不一次性载入所有图片,就是在定义dataset 类时 例如如下代码:init函数不把全部图片读取出来,而是通过另一个函数get_data 来载入,这样的的话每次迭代时都会重新打开文件获取相应data,我自己试了一下,我的机子的内存是16g,这个file 14g,用这个方法训练模型内存没有爆,但是这个方法比较耗cpu。这只是个人的拙见,还请其他大佬指正!

class PartPrimitive(gdata.Dataset):
"""
dataset for (part, block program) pairs
"""
def __init__(self, file_path):
    self.file_path = file_path
    self.get_data(0)
def __getitem__(self, index):
    return self.get_data(index)
def get_data(self,idx):
    f = h5py.File(self.file_path, 'r')
    data = f['data']
    labels = f['label']
    assert data.shape[0] == labels.shape[0]
    self.num = data.shape[0]
    data = data[idx]
    label = labels[idx, :, 0]
    param = labels[idx, :, 1:]
    data = data.astype(np.int64)
    label = label.astype(np.int64)
    param = param.astype(np.float32)
    return data,label,param

def __len__(self):
    return self.num
1赞

FCN那一章有代码:

class VOCSegDataset(gdata.Dataset):
    def __init__(self, is_train, crop_size, voc_dir, colormap2label):
        self.rgb_mean = nd.array([0.485, 0.456, 0.406])
        self.rgb_std = nd.array([0.229, 0.224, 0.225])
        self.crop_size = crop_size
        features, labels = read_voc_images(root=voc_dir, is_train=is_train)
        self.features = [self.normalize_image(feature) for feature in self.filter(features)]
        self.labels = self.filter(labels)
        self.colormap2label = colormap2label
        print('read ' + str(len(self.features)) + ' examples')

    def normalize_image(self, img):
        return (img.astype('float32') / 255 - self.rgb_mean) / self.rgb_std

    def filter(self, imgs):
        return [img for img in imgs if (
            img.shape[0] >= self.crop_size[0] and
            img.shape[1] >= self.crop_size[1])]

    def __getitem__(self, idx):
        feature, label = voc_rand_crop(self.features[idx], self.labels[idx], * self.crop_size)
        return feature.transpose((2, 0 ,1)), voc_label_indices(label, self.colormap2label)

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

然后通过DataLoader批量加载即可

voc_train = VOCSegDataset(True, crop_size, voc_dir, colormap2label)
voc_test = VOCSegDataset(False, crop_size, voc_dir, colormap2label)

batch_size = 64
num_workers = 0 if sys.platform.startswith('win') else 4
train_iter = gdata.DataLoader(voc_train, batch_size, shuffle=True, last_batch='discard', num_workers=num_workers)
test_iter = gdata.DataLoader(voc_test, batch_size, last_batch='discard', num_workers=num_workers)