title: Read TFRecord image data with new TensorFlow Dataset API
body: <p>I am having trouble reading TFRecord format image data using the "new" (TensorFlow v1.4) Dataset API. I believe the problem is that I am somehow consuming the whole dataset instead of a single batch when trying to read. I have a working example of doing this using the batch/file-queue API here: <a href="https://github.com/gnperdue/TFExperiments/tree/master/conv" rel="nofollow noreferrer">https://github.com/gnperdue/TFExperiments/tree/master/conv</a> (well, in the example I am running a classifier, but the code to read the TFRecord images is in the <code>DataReaders.py</code> class).</p>
  
  <p>The problem functions are, I believe, these:</p>
  
  <pre><code>def parse_mnist_tfrec(tfrecord, features_shape):
      tfrecord_features = tf.parse_single_example(
          tfrecord,
          features={
              'features': tf.FixedLenFeature([], tf.string),
              'targets': tf.FixedLenFeature([], tf.string)
          }
      )
      features = tf.decode_raw(tfrecord_features['features'], tf.uint8)
      features = tf.reshape(features, features_shape)
      features = tf.cast(features, tf.float32)
      targets = tf.decode_raw(tfrecord_features['targets'], tf.uint8)
      targets = tf.one_hot(indices=targets, depth=10, on_value=1, off_value=0)
      targets = tf.cast(targets, tf.float32)
      return features, targets
  
  class MNISTDataReaderDset:
      def __init__(self, data_reader_dict):
          # doesn't matter here
  
      def batch_generator(self, num_epochs=1):
          def parse_fn(tfrecord):
              return parse_mnist_tfrec(
                  tfrecord, self.name, self.features_shape
              )
          dataset = tf.data.TFRecordDataset(
              self.filenames_list, compression_type=self.compression_type
          )
          dataset = dataset.map(parse_fn)
          dataset = dataset.repeat(num_epochs)
          dataset = dataset.batch(self.batch_size)
          iterator = dataset.make_one_shot_iterator()
          batch_features, batch_labels = iterator.get_next()
          return batch_features, batch_labels
  </code></pre>
  
  <p>Then, in use:</p>
  
  <pre><code>        batch_features, batch_labels = \
              data_reader.batch_generator(num_epochs=1)
  
          sess.run(tf.local_variables_initializer())
          coord = tf.train.Coordinator()
          threads = tf.train.start_queue_runners(coord=coord)
          try:
              # look at 3 batches only
              for _ in range(3):
                  labels, feats = sess.run([
                      batch_labels, batch_features
                  ])
  </code></pre>
  
  <p>This generates an error like:</p>
  
  <pre><code> [[Node: Reshape_1 = Reshape[T=DT_UINT8, Tshape=DT_INT32](DecodeRaw_1, Reshape_1/shape)]]
   Input to reshape is a tensor with 50000 values, but the requested shape has 1
   [[Node: Reshape_1 = Reshape[T=DT_UINT8, Tshape=DT_INT32](DecodeRaw_1, Reshape_1/shape)]]
   [[Node: IteratorGetNext = IteratorGetNext[output_shapes=[[?,28,28,1], [?,10]], output_types=[DT_FLOAT, DT_FLOAT], _device="/job:localhost/replica:0/task:0/device:CPU:0"](OneShotIterator)]]
  </code></pre>
  
  <p>Does anyone have any ideas?</p>
  
  <p>I have a gist with the full code in the reader example and a link to the TFRecord files (our old, good friend MNIST, in TFRecord form) here:</p>
  
  <p><a href="https://gist.github.com/gnperdue/56092626d611ae23370a21fdeeb2abe8" rel="nofollow noreferrer">https://gist.github.com/gnperdue/56092626d611ae23370a21fdeeb2abe8</a></p>
  
  <p>Thanks!</p>
  
  <p>Edit - I also tried a <code>flat_map</code>, e.g.:</p>
  
  <pre><code>def batch_generator(self, num_epochs=1):
      """
      TODO - we can use placeholders for the list of file names and
      init with a feed_dict when we call `sess.run` - give this a
      try with one list for training and one for validation
      """
      def parse_fn(tfrecord):
          return parse_mnist_tfrec(
              tfrecord, self.name, self.features_shape
          )
      dataset = tf.data.Dataset.from_tensor_slices(self.filenames_list)
      dataset = dataset.flat_map(
          lambda filename: (
              tf.data.TFRecordDataset(
                  filename, compression_type=self.compression_type
              ).map(parse_fn).batch(self.batch_size)
          )
      )
      dataset = dataset.repeat(num_epochs)
      iterator = dataset.make_one_shot_iterator()
      batch_features, batch_labels = iterator.get_next()
      return batch_features, batch_labels
  </code></pre>
  
  <p>I also tried using just one file and not a list (in my first way of approaching this above). No matter what, it seems TF always wants to eat the entire file into the <code>TFRecordDataset</code> and won't operate on single records.</p>