如何使用 Dataset API 最大化张量流 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 使用率，您可以通过添加任何转换来添加数据并行性，并且您还可以考虑使用它来同时处理许多输入源并避免 I/O 阻塞。num_parallel_calls=NDataset.map()tf.contrib.data.parallel_interleave()

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