不使用 GPU 的 Tensorflow 数据集 API

Question

1. 问题：

我有一个tf.data.Dataset我用train_on_batch.

我的数据集如下所示：

Generate TFRecord path > tf.data.TFRecordDataset > Parse single example > Batch(2) > Map(merge) > Map(normalize) > Map(split to inputs,labels) > Batch(batch_size) > Prefetch(1)

我曾经RunMetadata用 Chrome 输出一个可读的时间轴。看起来IteratorGetNext只在 CPU 上运行并且占用了大量时间。

（我无法发布图片，IteratorGetNext耗时 617 毫秒，MEMCPYHtoD耗时 58 毫秒，训练耗时 500 毫秒）

我似乎找不到让 IteratorGetNext 在 GPU 上运行的方法，即使是部分运行。目前，CPU 使用率为 100%，GPU 使用率最高为 40-60%。

我希望是这样的：

Read from disk > Move from CPU to GPU > Preprocess.

我目前只使用一个 GPU，但我计划稍后使用更多 GPU，因此可扩展的解决方案将是完美的！

顺便说一下，我在带有 CUDA 10.0 和 python 3.6.7 的 Windows 10 上使用 tensorflow-gpu 1.13.1。我没有使用急切模式。我还没有在 Ubuntu 上试过，但这是一种可能性。

2.我试过的：

我尝试在管道中的几个地方使用prefetch_to_device和copy_to_devicefrom tf.data.experimental。

使用 时copy_to_device，IteratorGetNext 花费了两倍的时间。看起来它是在 GPU 上复制以仅复制回 CPU，因为MEMCPYHtoDIteratorGetNext 之后仍然存在。

我尝试将 Keras 替换为train_on_batch，session.run(train_op)但它并没有真正改善，我注意到的唯一变化是确实发生了一些预取，减少了几个样本的 IteratorGetNext 时间（与我在“预取”中放入的数量无关）。

对了，prefetch(1)还是prefetch(tf.data.experimental.AUTOTUNE)好像没有什么影响。

我尝试session.run了有和没有copy_to_device.

我还尝试将数据集的构建放在with tf.device("/gpu:0").

3.一些代码：

Generate TFRecord path > tf.data.TFRecordDataset > Parse single example > Batch(2) > Map(merge) > Map(normalize) > Map(split to inputs,labels) > Batch(batch_size) > Prefetch(1)

最后，我要补充一点，我的模型可能不够大，我可以通过让它“更大”来提高比率，但这并不是一个很好的解决方案。

- 编辑 ：

我有：

dataset = tf.data.Dataset.from_generator(self.random_shard_filepath_generator,
                                                 output_types=tf.string,
                                                 output_shapes=())

dataset = tf.data.TFRecordDataset(dataset)
dataset = dataset.map(lambda serialized_shard: self.parse_shard(serialized_shard, output_labels))

dataset = dataset.batch(self.shards_per_sample)
dataset = dataset.map(self.join_shards_randomly)
dataset = dataset.map(self.normalize_batch)
dataset = dataset.map(self.split_batch_io)

dataset = dataset.batch(batch_size).prefetch(1)

autoencoder.train_on_batch(dataset)

我改为：

...
dataset = dataset.batch(batch_size).prefetch(1)
autoencoder.train_on_batch(dataset)

感谢您EdoardoG让我尝试MultiDeviceIterator，这让我Iterator在 Keras 之外创建了一个train_on_batch.

现在IteratorGetNext只需要大约 0.05 毫秒，而以前大约需要 600 毫秒。

Answer 1

据我所知，Dataset API 操作通常在 CPU 上运行，因此无法在 GPU 上运行输入管道实际上是正常的。

有人编写了一个迭代器可以解决您的问题。