tensorflow - `tf.data.Dataset` 在 CPU 上运行，除了 `PrefetchDataset` ？

Question

在阅读了 tf.data 文档( here for TF 1.15 )、相关 TF 代码(Python 和 C++)后，我意识到其中大部分似乎纯粹在 CPU 上运行，除了 PrefetchDataset。

这是真的吗？

prefetch_to_device 的文档说:

NOTE: Although the transformation creates a tf.data.Dataset, the transformation must be the final Dataset in the input pipeline.

这表明所有其他数据集都无法处理此类基于 GPU 的数据集。

在查看代码时，似乎有一些内部数据集，例如_CopyToDeviceDataset 和 _MapOnGpuDataset，可能处理 GPU 数据集。

如果我想在 GPU 上进行预处理(例如数据增强、其他一些巧妙的非平凡的东西)，这意味着我不能使用 tf.data ？ (我也想使用图形模式，但不确定这是否相关。)

现在我还找到了_GeneratorDataset。 That kernel also is defined on GPU 。这意味着如果我的 next_func 在 GPU 上返回张量，它将始终保留在 GPU 上？

Answer 1

你是对的，tf.data 目前在 CPU 上进行所有处理。通常，这是为了避免与 GPU 竞争，并且因为许多预处理在 CPU 上更容易实现。然而，没有什么从根本上阻止 tf.data 在 GPU 上进行处理 - 这只是实现此类支持的问题。从表面上看，map_on_gpu确实提供了一种在 GPU 上应用 map 函数的方法，尽管 map_on_gpu 尚未在公共(public) API 中导出。如果您对此类功能感兴趣，请创建 Github Issue描述您的用例

为了回答你的最后一个问题，如果你有一个放置在 GPU 上的 _GeneratorDataset，它生成的张量也将位于 GPU 上。但是，您无法应用除预取之外的其他数据集转换，因为这些数据集转换没有 GPU 实现。