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

Question

shuffle对于中等缓冲区大小(例如 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))从磁盘读取。像物理长期存储，可能是磁性旋转硬盘驱动器。这很慢。如果您有能力将所有这些存储在 ram 中，或者存储在超快的 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 行开始并尝试在长度为 1000 的缓冲区中打乱您的 32 批次。每次训练循环请求新的训练批次时都会发生这种情况。洗牌步骤将所有这些大批量洗牌，挑选出第一个并添加一个新的……每……一次……时间。
我们可以像这样颠倒批次和洗牌的顺序

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

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