title: Epoch counter with TensorFlow Dataset API
body: <p>I'm changing my TensorFlow code from the old queue interface to the new <a href="https://www.tensorflow.org/api_docs/python/tf/data/Dataset" rel="noreferrer">Dataset API</a>. In my old code I kept track of the epoch count by incrementing a <code>tf.Variable</code> every time a new input tensor is accessed and processed in the queue. I'd like to have this epoch count with the new Dataset API, but I'm having some trouble making it work.</p>
  
  <p>Since I'm producing a variable amount of data items in the pre-processing stage, it is not a simple matter of incrementing a (Python) counter in the training loop - I need to compute the epoch count with respect to the input of the queues or Dataset.</p>
  
  <p>I mimicked what I had before with the old queue system, and here is what I ended up with for the Dataset API (simplified example):</p>
  
  <pre><code>with tf.Graph().as_default():
  
      data = tf.ones(shape=(10, 512), dtype=tf.float32, name="data")
      input_tensors = (data,)
  
      epoch_counter = tf.Variable(initial_value=0.0, dtype=tf.float32,
                                  trainable=False)
  
      def pre_processing_func(data_):
          data_size = tf.constant(0.1, dtype=tf.float32)
          epoch_counter_op = tf.assign_add(epoch_counter, data_size)
          with tf.control_dependencies([epoch_counter_op]):
              # normally I would do data-augmentation here
              results = (tf.expand_dims(data_, axis=0),)
              return tf.data.Dataset.from_tensor_slices(results)
  
      dataset_source = tf.data.Dataset.from_tensor_slices(input_tensors)
      dataset = dataset_source.flat_map(pre_processing_func)
      dataset = dataset.repeat()
      # ... do something with 'dataset' and print
      # the value of 'epoch_counter' every once a while
  </code></pre>
  
  <p>However, this doesn't work. It crashes with a cryptic error message:</p>
  
  <pre><code> TypeError: In op 'AssignAdd', input types ([tf.float32, tf.float32])
   are not compatible with expected types ([tf.float32_ref, tf.float32])
  </code></pre>
  
  <p>Closer inspection shows that the <code>epoch_counter</code> variable might not be accessible within the <code>pre_processing_func</code> at all. Does it live in a different graph perhaps?</p>
  
  <p>Any idea how to fix the above example? Or how to get the epoch counter (with decimal points, e.g. 0.4 or 2.9) through some other means?</p>