TensorFlow数据集API不使用GPU

Question

1. 问题:

我有一个 tf.data.Dataset 数据集，我将其传递给 Keras 模型 (tf.python.keras) 的 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%。

我希望能够像这样：

从磁盘读取 > 将数据从 CPU 移动到 GPU > 预处理。

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

顺便说一下，我正在使用 Tensorflow-GPU 1.13.1，在 Windows 10 上使用 CUDA 10.0 和 Python 3.6.7。我没有使用 eager 模式。我还没有在 Ubuntu 上尝试过，但这是一个可能性。

2. 我尝试了什么：

我尝试在管道中的几个位置使用 tf.data.experimental 中的 prefetch_to_device 和 copy_to_device。

使用 copy_to_device 时，IteratorGetNext 的时间变成了两倍。看起来它是在 GPU 上进行复制，然后再回到 CPU，因为 IteratorGetNext 之后仍然存在 MEMCPYHtoD。

我尝试使用 session.run(train_op) 替换 Keras 的 train_on_batch，但实际上并没有改善。我注意到的唯一变化是实际上进行了一些预取操作，将 IteratorGetNext 的时间减少了一些样本（与我放入 "prefetch" 中的数量无关）。

顺便说一下，prefetch(1) 或 prefetch(tf.data.experimental.AUTOTUNE) 似乎没有任何影响。

我尝试过 session.run 以及不带 copy_to_device 的 session.run。

我还尝试将数据集的构建放在 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 毫秒。

- Zelgunn

Answer 1

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

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

- EdoG

---

谢谢指出这个问题。 我看到了，但是我认为除非我使用多个GPU，否则它不会有用。 所以我尝试了一下，它起作用了，但实际上并不是由于这个迭代器。 它让我在Keras的train_function之外创建和初始化迭代器，并且它奏效了。所以我尝试使用make_initializable_iterator创建一个迭代器，它也奏效了。 - Zelgunn

我习惯于编写普通的TF代码而不是使用Keras，因为它对于训练GANs更加灵活，并且可以更好地理解发生了什么，特别是在train()函数和输入管道中。我无法确定您的加速是由于TF Dataset API还是GPU迭代器所致，但我经常遇到性能不稳定的情况，尤其是在使用多个数据集时切换一次性/可初始化/可喂养迭代器。我还提出了一个长问题（仍未得到答复），也许我应该尝试使用自定义迭代器。 - EdoG

Answer 2

使用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官网提供的一些有用指南：

- 20r25g22e

---

我尝试首先将我的main()包装在with tf.device("/gpu:0")中，然后构建模型，最后是主要的训练循环。但这要么没有改变太多，要么使情况更糟（IteratorGetNext花费的时间增加了3倍）。至于多GPU，我知道这种技术，我更关心的是将数据分配到多个GPU上的分布问题。如果CPU无法为1个GPU提供支持，那么对于更多的GPU也是如此。 - Zelgunn