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 毫秒。
train_function
之外创建和初始化迭代器,并且它奏效了。所以我尝试使用make_initializable_iterator
创建一个迭代器,它也奏效了。 - Zelgunn