memmap 数组到 pytorch 和梯度累积

Question

我有一个大型数据集（> 62 GiB）在处理后保存为两个 NumPy.memmap 数组，一个数据和另一个数据集具有这些形状的标签（7390,60,224,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,224,224,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 很容易，但我找不到如何处理 memmap 数组，就像在带有生成器的 tensorflow 中那样。

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

或者如果有人在 TF2

上有 GA 的现成代码

Answer 1

我只解决随机播放问题。

不要使用 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 已经内置了处理此类内容的方法

memmap 数组到 pytorch 和梯度累积

memmap arrays to pytorch and gradient accumulation

dataset

numpy-memmap

tensorflow

pytorch

batchsize