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

【问题标题】：Why is TensorFlow's tf.data.Dataset.shuffle so slow?为什么 TensorFlow 的 tf.data.Dataset.shuffle 这么慢？
【发布时间】：2018-01-13 13:59:11
【问题描述】：

以下代码中的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 是否除了简单地打乱数据之外还做了其他事情。

【问题讨论】：

你对慢的定义是什么？同样在您的 TensorFlow 代码中加载文件，而您的 numpy 示例没有。也可能是相关的：github.com/tensorflow/tensorflow/issues/14857
我没有计时，但 TF 的版本明显慢了很多。我的 numpy 块中的最后一行用于将文件放入内存并将它们转换为 TF 格式。但是谢谢链接。
在洗牌后尝试使用预取选项。它应该使加载更快。例如：dataset.prefetch(buffer_size=10)
我的理解是，TensorFlow shuffle 实际上会复制元素，如果你没有足够的 RAM 来保存它们，这些元素会变得非常缓慢。
2019 年，tf 1.14 cuda 10，训练集 84519 个样本，调整大小（224,224,3 from 512 x694）作为预处理中的唯一操作，批量大小为 16 tf.data.Dataset。 shuffle(len(train_files)) 在 GTX 1080 8GB GPU、Intel core i7 7820 HK 8 核 CPU 中需要 1 小时 5 分钟，开始准备数据需要大约 1 小时。

标签： tensorflow

【解决方案1】：

正如 Anton Codes 所写，您的第一个 sn-p 对从您的文件（可能是特征数据）解析的任何 _parse_function 进行洗牌，而您的第二个 sn-p 只洗牌文件名。

如果文件级别的改组就足够了，您实际上可以通过tf.data.Dataset API 实现（大致）相同的性能：

dataset = tf.data.Dataset.from_tensor_slices((filenames, labels))
dataset = dataset.shuffle(len(filenames)) # shuffle file names
dataset = dataset.map(_parse_function)
dataset = dataset.batch(batch_size)

这种将“指针”改组到训练样本而不是样本本身的做法通常可以提高性能。

NumPy 可能仍然会更高效一些，因为计算图内部的混洗开销（tf.data.Dataset.shuffle 确实如此，实际上有一个专门用于此操作的 C++ 内核）。

tf.data.Dataset 方法的优点是它可以在每个 epoch 之后自动重新洗牌。

【讨论】：

【解决方案2】：

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

您的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 个批次进行洗牌。每次训练循环请求一个新的训练批次时都会发生这种情况。洗牌步骤洗牌所有这些大批量，挑选出第一个并添加一个新的......每......一次......时间。

我们可以像这样颠倒batch和shuffle的顺序

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 似乎几乎没有用，除了项目是一次随机从批次中拉出的，因此实际上有机会不被绘制并与以后的数据混合。只是没有那么多。

【讨论】：