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花了617ms，MEMCPYHtoD花了58ms，训练花了500ms）

我似乎找不到找到使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_device和tf.data.experimental。

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

我尝试用train_on_batch替换Keras的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的Iterator之外创建了train_on_batch。

现在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可以解决您的问题。

Tensorflow Dataset API未使用GPU

2 个答案: