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%。在

我希望有这样的东西：

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

我目前只使用一个GPU，但我计划以后使用更多的GPU，这样一个可伸缩的解决方案将是完美的！在

顺便说一下，我在Windows10上使用了TensorFlowGPU1.13.1和CUDA10.0和Python3.6.7。我没有使用急切模式。 我还没试过Ubuntu，但有可能。在

2。我尝试了：

我尝试在管道中的几个位置使用prefetch_to_device和{}，它们来自{}。在

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

我尝试用session.run(train_op)替换Keras'train_on_batch，但它并没有真正改善，我注意到的唯一变化是实际发生了一些预取，减少了一些示例的iteratorGet（与我在“prefetch”中投入的量无关）。在

顺便说一下，prefetch(1)或{}似乎没有任何影响。在

我尝试了session.run，无论有没有copy_to_device。在

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

3。一些代码：

^{pr2}$

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

--编辑：

我有：

...
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.05ms，而以前只需要600毫秒

Answer 1

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

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

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's website的一些有用指南