为什么TensorFlow的tf.data.Dataset.shuffle这么慢？

Question

shuffle对于中等的buffer_size（例如1000），以下代码中的步骤非常慢：

filenames = tf.constant(filenames)
dataset = tf.data.Dataset.from_tensor_slices((filenames, labels))
dataset = dataset.map(_parse_function)
dataset = dataset.batch(batch_size)
dataset = dataset.shuffle(buffer_size)

如果我们使用numpy随机数据，则代码如下所示：

idx = np.arange(len(filenames))
np.random.shuffle(idx)
new_filenames = [filenames[i] for i in idx]
next_batch_filenames = new_filenames[:batch_size]
# get the corresponding files in batch

这要快得多。我想知道TF是否可以做一些超出随机整理数据的事情。

Answer 1

比较的是两种完全不同的操作。

您dataset = tf.data.Dataset.from_tensor_slices((filenames, labels))从磁盘读取。就像物理长期存储一样，可能是磁旋转硬盘驱动器。这很慢。如果您有能力将所有这些存储在内存中，或者存储在超快的 raid 式闪存驱动器上，那么您将解决最大的瓶颈。

_parse_function每次读取数据时，您还会为每个数据点触发一个。该解析的计算将需要时间，并且根据其中的内容，它可能很重要。

与 numpy 的比较并不公平，因为您的 numpy 示例不涉及从磁盘读取或解析数据。

这应该是差异的主要部分。如果您已经解决了上述问题，那么下一个寻求更高加速的地方就是这些行

3) dataset = dataset.map(_parse_function)
4) dataset = dataset.batch(batch_size)
5) dataset = dataset.shuffle(buffer_size)

这些是您的代码行。第 4 行生成批量数据，可能是 32 个（batch_size当然）。然后第 5 行启动并尝试将 32 个批次放入长度为 1000 的缓冲区中。每次训练循环请求新的训练批次时都会发生这种情况。洗牌步骤对所有这些大批次进行洗牌，挑选出第一个批次并添加一个新批次......每次......一次。

我们可以像这样颠倒批处理和洗牌的顺序

3) dataset = dataset.map(_parse_function)
4) dataset = dataset.shuffle(buffer_size)
5) dataset = dataset.batch(batch_size)

无论如何，这样更好，因为以前批次的内容总是相同的，但顺序是混合的。这样批次的内容也会被随机化。接下来，洗牌只能洗牌 1000 个项目，而不是 32x1000 个项目。最后，我们可以质疑是否真的需要 1000 的缓冲区大小。假设我们的数据集有 2000 个项目。缓冲区大小为 320 且批次大小为 32 肯定会很好地随机化我们的数据，有效地为缓冲区中的任何数据提供 10% 的进入下一个批次的数据，以及 90% 的数据被推回与其他数据混合的数据。那很好。缓冲区大小为 64 且批处理大小为 64 似乎几乎没有用，除了一次随机地从批处理中拉出一个项目之外，因此实际上有可能不被绘制并与后面的数据混合。只是没那么多。