memmap 数组到 pytorch 和梯度累积

问题描述

处理后我有一个大数据集 (> 62 GiB) 保存为两个 NumPy.memmap 数组,一个数据,另一个用于标签,数据集具有这些形状 (7390,60,224,3) 和 (7390) 和洗牌,所以我需要先洗牌。

现在我使用 tensorflow2 并在我的 generator 之前使用此代码来管理 memmap 数组

def my_generator():
    for i in range(len(numpy_array)):
          yield numpy_array[i,:,:],np.array(labels[i]).reshape(1)

full_dataset = tf.data.Dataset.from_generator(
    generator=my_generator,output_types=(np.uint8,np.int32),output_shapes=((60,3),(1))
)

full_dataset = full_dataset.shuffle(SHUFFLE_BUFFER_SIZE,reshuffle_each_iteration=False)
train_dataset = full_dataset.take(train_size)
test_dataset = full_dataset.skip(train_size)
val_dataset = test_dataset.skip(test_size)
test_dataset = test_dataset.take(test_size) 

这样我就可以训练而无需通过混洗和批处理将整个数据集加载到内存中。

现在使用当前的模型和数据集,vram 不足以加载超过 2 个批次作为张量。 我不能用 2 的批量进行训练。

我想到了梯度累积,但我不能用 TF2 做到这一点,我发现用 pytorch 很容易,但我找不到如何像在带有生成器的 tensorflow 中那样使用 shuffle 和 split 处理 memmap 数组。

所以我需要知道如何在 pytorch 中使用相同的改组和批处理从 pytorch 加载数据集。

或者如果有人有 TF2 上的 GA 的现成代码

解决方法

我将只解决随机播放问题。

不要使用 tf.data.Dataset 进行改组,而是在生成器级别进行。这应该有效:

class Generator(object):
    def __init__(self,images,labels,batch_size):
        self.images = images
        self.labels = labels
        self.batch_size  = batch_size
        self.idxs = np.arange(len(self.images))
        self.on_epoch_end()

    def on_epoch_end(self):
        # Shuffle the indices
        np.random.shuffle(self.idxs)

    def generator(self):
        i = 0
        while i < len(self.idxs):
            idx = self.idxs[i]
            yield (self.images[idx],self.labels[i])
            i += 1
        self.on_epoch_end()

    def batch_generator(self):
        it = iter(self.generator)
        while True:
            vals = [next(it) for i in range(self.batch_size)]
            images,labels = zip(*vals)
            yield images,labels

然后你就可以使用它了

gen = Generator(...)
it = iter(gen)

batch = next(it)  # Call this every time you want a new batch

我确信 pytorch 已经为这类东西内置了方法