tensorflow - 如何使用 Dataset API 最大化 tensorflow GPU 训练中的 CPU 利用率？

Question

我正在使用新的数据集 API 来训练简单的前馈 DL 模型。我对最大化训练速度感兴趣。由于我的网络规模不大，正如预期的那样，我看到 GPU 利用率较低。没事儿。但我不明白的是为什么CPU使用率也远未达到100%。我使用的是多 cpu/gpu 核心机器。目前，batch_size = 128 时，我的速度可达 140 步/秒。如果我缓存数据集，我最多可达 210 步(初始扫描后)。所以我希望通过足够的预取，我应该能够在没有缓存的情况下达到相同的速度。然而，使用各种预取和 prefetch_to_device 参数，我无法获得超过 140 步/秒的速度。我还将 num_parallel_calls 设置为 cpu 核心数，大约提高了 20%。

理想情况下，我希望预取线程位于与输入管道的其余部分不相交的CPU核心上，这样它提供的任何好处都是严格相加的。但从 CPU 使用情况分析来看，我怀疑预取和输入处理发生在每个核心上:

有没有办法更好地控制CPU分配？我尝试过 prefetch(1)、prefetch(500) 和其他几个值(在批处理之后或数据集构建结束时)，以及与 prefetch_to_device(gpu_device, batch_size = None, 1, 500 等) 结合使用)。到目前为止，不带 prefetch_to_device 的 prefetch(500) 效果最好。

为什么预取不会尝试耗尽我机器上的所有 CPU 功率？训练速度还有哪些其他可能的瓶颈？

非常感谢!

Answer 1

Dataset.prefetch(buffer_size)转换为您的输入管道添加了管道并行性和(有界)缓冲。因此，增加 buffer_size 可能会增加 Dataset.prefetch() 输入运行的时间(因为缓冲区更有可能有可用空间)，但它不会提高输入运行的速度(从而提高 CPU 使用率)。

通常，要提高管道速度并增加 CPU 使用率，您可以通过将 num_parallel_calls=N 添加到任意 Dataset.map() 来添加数据并行性。转换，您也可以考虑使用 tf.contrib.data.parallel_interleave()同时处理多个输入源并避免 I/O 阻塞。

tf.data Performance Guide有关如何提高输入管道性能的更多详细信息，包括这些建议。