如何在'input_fn'中使用张量流的迭代器'make_initializable_iterator'？

Question

I want to train my mode with a tf.estimator.Estimator and load my data by Dataset API.Because my data,for example 'mnist', is a array(tensor),so I try to load it with 'tf.data.Dataset.from_tensor_slices'.But I don't how to initialize 'make_initializable_iterator' within a 'input_fn'.

If I can use 'make_one_shot_iterator' to train successfully, but it load slowly before training. And 《 Higher-Level APIs in TensorFlow 》is a good example to 'make_initializable_iterator' within a 'input_fn',but it needs to return a 'iterator_initializer_hook' to other function from 'input_fn' . I want to know is there any other better or more elegant way?

    def input_fn():

    mnist_data = input_data.read_data_sets('mnist_data', one_hot=False)
    images = mnist_data.train.images.reshape([-1, 28, 28, 1])
    labels = np.asarray(mnist_data.train.labels, dtype=np.int64)

    # Build dataset iterator
    dataset = tf.data.Dataset.from_tensor_slices((images, labels))
    dataset = dataset.repeat(None)  # Infinite iterations
    dataset = dataset.shuffle(buffer_size=10000)
    dataset = dataset.batch(100)
    iterator = dataset.make_one_shot_iterator()
    next_example = iterator.get_next()
    # Set runhook to initialize iterator

    return next_example

Answer 1

In TensorFlow version 1.5 and later, the tf.estimator.Estimator will automatically create and initialize an initializable iterator when you return a tf.data.Dataset from your input_fn . This enables you to write the following code, without having to worry about initialization or hooks:

def input_fn():
    mnist_data = input_data.read_data_sets('mnist_data', one_hot=False)
    images = mnist_data.train.images.reshape([-1, 28, 28, 1])
    labels = np.asarray(mnist_data.train.labels, dtype=np.int64)

    # Build dataset.
    dataset = tf.data.Dataset.from_tensor_slices((images, labels))
    dataset = dataset.repeat(None)  # Infinite iterations
    dataset = dataset.shuffle(buffer_size=10000)
    dataset = dataset.batch(100)
    return dataset

Answer 2

Inside your code, add this:

      self.hooks.append(utils_hooks.DatasetHook(iter))

In the run_loop.py, before the call into your fn, add this

 for hook in dataset_hooks:
        sess.run(hook.iterator().initializer)

Then, it should be fine.