How to use tf.data's initializable iterators within a tf.estimator's input_fn?

Question

I would like to manage my training with a tf.estimator.Estimator but have some trouble to use it alongside the tf.data API.

I have something like this:

Answer 1

As of TensorFlow 1.5, it is possible to make input_fn return a tf.data.Dataset, e.g.:

def input_fn():
  dataset = tf.data.TextLineDataset("test.txt")
  # map, shuffle, padded_batch, etc.
  return dataset

See c294fcfd.

---

For previous versions, you can add the iterator's initializer in the tf.GraphKeys.TABLE_INITIALIZERS collections and rely on the default initializer.

tf.add_to_collection(tf.GraphKeys.TABLE_INITIALIZERS, iterator.initializer)