Reinitializable iterators for simultaneous training and validation

Question

I want to use Dataset and Iterators to evaluate on a validation set during training. I want to evaluate on one (or a few) validation batches every now and then — that every now and then is typically not an epoch.

However reinitializable iterators start all over again when reinitalized to switch their input. E.g.

import tensorflow as tf

dataset_trn = tf.data.Dataset.range(10)
dataset_tst = tf.data.Dataset.range(10).map(lambda i: i + 1000)
iterator = tf.data.Iterator.from_structure(dataset_trn.output_types,
    dataset_trn.output_shapes)
batch = iterator.get_next()
trn_init_op = iterator.make_initializer(dataset_trn)
tst_init_op = iterator.make_initializer(dataset_tst)

sess = tf.InteractiveSession()

for _ in range(2):
  sess.run(trn_init_op)
  for _ in range(5):
    print(batch.eval())

  sess.run(tst_init_op)
  print(batch.eval())

returns

0
1
2
3
4
1000
0
1
2
3
4
1000

but I would like it to resume training like that:

0
1
2
3
4
1000
5
6
7
8
9
1001

Is there a way to achieve this? Note that in practice, batches are shuffled, and I would like it to resume at the same pseudo-random point.

Answer 1

Feedable iterators should help, but they're tough to work with. You need to create a placeholder and string handles:

dataset_trn = tf.data.Dataset.range(10)
dataset_tst = tf.data.Dataset.range(10).map(lambda i: i + 1000)

holder = tf.placeholder(tf.string, shape=[])
iterator = tf.data.Iterator.from_string_handle(
    holder, dataset_trn.output_types, dataset_trn.output_shapes)
batch = iterator.get_next()

trn_iter = dataset_trn.make_one_shot_iterator()
trn_handle = trn_iter.string_handle()

tst_iter = dataset_tst.make_one_shot_iterator()
tst_handle = tst_iter.string_handle()

with tf.Session() as sess:
  for _ in range(2):

    trn_string = sess.run(trn_handle)
    tst_string = sess.run(tst_handle)

    for _ in range(5):
      print(sess.run(batch, feed_dict={holder: trn_string}))

    print(sess.run(batch, feed_dict={holder: tst_string}))