python - 如何将 prefetch_to_device 与 make_initializer 结合使用

Question

根据官方文档， Tensorflow 1.8已经引入tf.contrib.data.prefetch_to_device，可以配合使用。tf.data.Dataset.apply

我已经编写了一些自定义数据加载，我想利用它来预取到 GPU。问题是这个操作需要是数据集所有转换中的最后一个。我正在使用tf.data.Iterator.make_initializer，tf.data.Iterator.from_structure如下所示。这是引发错误的最小示例。

import tensorflow as tf

class MyData(object):
    def __call__(self):
         return range(100)

expected_shapes = []
expected_types = tf.int32
iterator = tf.data.Iterator.from_structure(output_types=expected_types, output_shapes=expected_shapes)
dataset = tf.data.Dataset.from_generator(MyData(), output_types=expected_types, output_shapes=expected_shapes)

prefetch_op = tf.contrib.data.prefetch_to_device(device="/gpu:0")
dataset = dataset.apply(prefetch_op)

# This raises NotImplementedError: `prefetch_to_device()` must be the last transformation in a dataset pipeline.
initializer = iterator.make_initializer(dataset)

当然，当我将顺序翻转时不会发生这种情况，但是初始化器不会考虑预取转换。

initializer = iterator.make_initializer(dataset)
dataset = dataset.apply(prefetch_op)

在告诉它预取之后，我怎样才能让它工作来创建初始化程序？或者这根本不可能（目前）？与来自 tensorflow 的这个测试用例Iterator不同，我想避免创建一个新的.

标签： pythontensorflowprefetch