memmap阵列到pytorch和梯度积累

Question

我有一个大数据集(>62 GiB(，在处理后保存为两个NumPy.memmap数组，一个是数据，另一个是标签，数据集有这些形状(7390，62022422424，3(和(7390(，并且是NOT混洗的，所以我需要先混洗它。

现在我使用tensorflow2，并将此代码与生成器一起使用，以在之前管理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很容易，但我找不到如何像用生成器处理tensorflow中那样用shuffle和split处理memmap数组。

所以我需要知道如何用pytorch中相同的混洗和批处理从pytorch加载基准集。

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

Answer 1

我只讨论shuffle问题。

不要使用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已经通过内置了这种东西的方法