TensorFlow Dataset API 不使用 GPU

Question

1.问题：

我有一个 tf.data.Dataset，我用 train_on_batch 将它提供给 Keras 模型 (tf.python.keras)。

我的数据集如下所示：

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。我尝试了什么：

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

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

我尝试将 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。一些代码：

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)

我改成了：

...
dataset = dataset.batch(batch_size).prefetch(1)
dataset_iterator = dataset.make_initializable_iterator()
dataset_initializer = dataset_iterator.initializer

session.run(dataset_initializer)

x, y = dataset_iterator
autoencoder.train_on_batch(x, y)

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

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

Answer 1

使用with tf.device('/gpu:0'): 包装您的NN 代码，其中gpu:0 是您系统中的第一个gpu。

如果您想使用多个 GPU：

for d in ['/device:GPU:2', '/device:GPU:3']:
  with tf.device(d):
    <your code here>

来自tensorflow's website的一些有用指南

Answer 2

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

有人写了an iterator 可以解决您的问题。