Tensorflow tf.dataset.shuffle 非常慢

Question

我正在用 9100 张图像（每个尺寸为 256 x 64）训练 VAE model。 我用 Nvidia RTX 3080 训练 model。首先，我将所有图像加载到大小为 9100 x 256 x 64 的 numpy 数组中，称为traindata 。 然后，为了形成一个训练数据集，我使用

train_dataset = (tf.data.Dataset.from_tensor_slices(traindata).shuffle(len(traindata)).batch(batch_size))

这里我使用了65的batch_size 。关于训练时看到的东西，我主要有2个问题：

问题一：

根据文档，整个数据集在每个时期都被重新洗牌。 但是，这种方式的训练非常慢（每个 epoch 大约 50 秒）。 我通过在创建数据集时不调用.shuffle(len(traindata))与没有 shuffle 的训练进行了比较，训练速度要快得多（大约 20s/epoch）。 我想知道为什么.shuffle()操作如此缓慢，是否有任何方法可以使其更快？ 根据这个 StatsSE thread ，洗牌对于训练非常重要，这就是我包含洗牌操作的原因。

问题2：

当我在创建数据集时调用.shuffle()时，Tensorflow 总是给出以下消息

I tensorflow/core/platform/windows/subprocess.cc:308] SubProcess ended with return code: 4294967295

我尝试在网上搜索，但仍然无法理解这背后的含义。 这是否意味着有一些错误，或者这只是一个我可以忽略的警告？

Answer 1

这是因为将数据集的所有元素保存在缓冲区中是昂贵的。 除非您绝对需要完美的随机性，否则您应该使用较小的buffer_size 。 最终将采用所有元素，但以更具确定性的方式。

这就是使用较小的buffer_size会发生的事情，比如 3。缓冲区是括号，Tensorflow 在这个括号中采样一个随机值。 随机抽取的是^

1) [1 2 3]4 5 6 7 8 9 
      ^
2) [1 3 4]5 6 7 8
        ^
3) [1 3 5]6 7 8
        ^
4) [1 3 6]7 8
    ^
5) [3 6 7]8

ETC

因此，较早的值将在您的 epoch 中较早地获取，但您仍将完成一些洗牌，并且最终将获取所有样本。

tl;博士将buffer_size减少了很多