title: How do I iterate explicitly over batches in TensorFlow dataset?
body: <p>I have a Computer Vision application that analyzes very large (30k x 40x) images by breaking them into tiles; this tiling process produces a TensorFlow TFRecord file with one TFRecord per tile, potentially a large number like 5000.   What I would like to explicitly iterate over all the TFRecord in batches, and for each bunch, I want to use a Deep Learning model to produce a segmentation map for each tile.  So I would have something like:</p>
  <pre><code>    for input_tiles_tensor in break_into_batches( TFRecord_file, batch_size ):
          output_tiles_tensor is model.predict( input_tiles_tensors )
          ...use output_tile_tensor to generate output, written to file...
  </code></pre>
  <p>The part I need help with is <code>break_into_batches</code>   Suggestions welcome!</p>
  <p>-------------------------------- ADDITION ---------------------------------------</p>
  <p>I think I see an approach, BUT...   Here's the approach:</p>
  <pre><code>AUTO = tf.data.experimental.AUTOTUNE
  image_feature = {
      'image': tf.io.FixedLenFeature([], tf.string),
  }
  def _parse_image(example_proto):
      example = tf.io.parse_single_example(example_proto, image_feature, name = &quot;parse_example&quot;)
      dim_with_overlap = P[ 'DIM' ] + 2 * P[ 'PIXEL_OVERLAP' ]
      image = tf.reshape( tf.io.decode_raw(example['image'],out_type=np.dtype('uint8')), 
                          (dim_with_overlap, dim_with_overlap, 3))
      return image
  
  
  dataset = tf.data.TFRecordDataset( &quot;../input/blortzk-hubmap-test-overlapping-tiled-images/test/26dc41664-6068.tfrec&quot;,
                                     compression_type = &quot;GZIP&quot; )
  dataset = dataset.batch( 77 )
  dataset = dataset.prefetch(AUTO)
  dataset = dataset.map(_parse_image)
  
  iterator = iter( dataset )
  batch = iterator.get_next()
  print( &quot;type( batch )&quot;, type( batch ), &quot;batch.shape&quot;, batch.shape, file = sys.stderr )
  </code></pre>
  <p>The BUT is that this gives an error in the <code>dataset.map</code> statement:</p>
  <pre><code>---------------------------------------------------------------------------
  ValueError                                Traceback (most recent call last)
  &lt;ipython-input-25-67a2e02c344d&gt; in &lt;module&gt;
       15 dataset = dataset.batch( 77 )
       16 dataset = dataset.prefetch(AUTO)
  ---&gt; 17 dataset = dataset.map(_parse_image)
       18 
       19 iterator = iter( dataset )
  
  /opt/conda/lib/python3.7/site-packages/tensorflow/python/data/ops/dataset_ops.py in map(self, map_func, num_parallel_calls, deterministic)
     1693     &quot;&quot;&quot;
     1694     if num_parallel_calls is None:
  -&gt; 1695       return MapDataset(self, map_func, preserve_cardinality=True)
     1696     else:
     1697       return ParallelMapDataset(
  
  /opt/conda/lib/python3.7/site-packages/tensorflow/python/data/ops/dataset_ops.py in __init__(self, input_dataset, map_func, use_inter_op_parallelism, preserve_cardinality, use_legacy_function)
     4043         self._transformation_name(),
     4044         dataset=input_dataset,
  -&gt; 4045         use_legacy_function=use_legacy_function)
     4046     variant_tensor = gen_dataset_ops.map_dataset(
     4047         input_dataset._variant_tensor,  # pylint: disable=protected-access
  
  /opt/conda/lib/python3.7/site-packages/tensorflow/python/data/ops/dataset_ops.py in __init__(self, func, transformation_name, dataset, input_classes, input_shapes, input_types, input_structure, add_to_graph, use_legacy_function, defun_kwargs)
     3369       with tracking.resource_tracker_scope(resource_tracker):
     3370         # TODO(b/141462134): Switch to using garbage collection.
  -&gt; 3371         self._function = wrapper_fn.get_concrete_function()
     3372         if add_to_graph:
     3373           self._function.add_to_graph(ops.get_default_graph())
  
  /opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/function.py in get_concrete_function(self, *args, **kwargs)
     2937     &quot;&quot;&quot;
     2938     graph_function = self._get_concrete_function_garbage_collected(
  -&gt; 2939         *args, **kwargs)
     2940     graph_function._garbage_collector.release()  # pylint: disable=protected-access
     2941     return graph_function
  
  /opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/function.py in _get_concrete_function_garbage_collected(self, *args, **kwargs)
     2904       args, kwargs = None, None
     2905     with self._lock:
  -&gt; 2906       graph_function, args, kwargs = self._maybe_define_function(args, kwargs)
     2907       seen_names = set()
     2908       captured = object_identity.ObjectIdentitySet(
  
  /opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/function.py in _maybe_define_function(self, args, kwargs)
     3211 
     3212       self._function_cache.missed.add(call_context_key)
  -&gt; 3213       graph_function = self._create_graph_function(args, kwargs)
     3214       self._function_cache.primary[cache_key] = graph_function
     3215       return graph_function, args, kwargs
  
  /opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/function.py in _create_graph_function(self, args, kwargs, override_flat_arg_shapes)
     3073             arg_names=arg_names,
     3074             override_flat_arg_shapes=override_flat_arg_shapes,
  -&gt; 3075             capture_by_value=self._capture_by_value),
     3076         self._function_attributes,
     3077         function_spec=self.function_spec,
  
  /opt/conda/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py in func_graph_from_py_func(name, python_func, args, kwargs, signature, func_graph, autograph, autograph_options, add_control_dependencies, arg_names, op_return_value, collections, capture_by_value, override_flat_arg_shapes)
      984         _, original_func = tf_decorator.unwrap(python_func)
      985 
  --&gt; 986       func_outputs = python_func(*func_args, **func_kwargs)
      987 
      988       # invariant: `func_outputs` contains only Tensors, CompositeTensors,
  
  /opt/conda/lib/python3.7/site-packages/tensorflow/python/data/ops/dataset_ops.py in wrapper_fn(*args)
     3362           attributes=defun_kwargs)
     3363       def wrapper_fn(*args):  # pylint: disable=missing-docstring
  -&gt; 3364         ret = _wrapper_helper(*args)
     3365         ret = structure.to_tensor_list(self._output_structure, ret)
     3366         return [ops.convert_to_tensor(t) for t in ret]
  
  /opt/conda/lib/python3.7/site-packages/tensorflow/python/data/ops/dataset_ops.py in _wrapper_helper(*args)
     3297         nested_args = (nested_args,)
     3298 
  -&gt; 3299       ret = autograph.tf_convert(func, ag_ctx)(*nested_args)
     3300       # If `func` returns a list of tensors, `nest.flatten()` and
     3301       # `ops.convert_to_tensor()` would conspire to attempt to stack
  
  /opt/conda/lib/python3.7/site-packages/tensorflow/python/autograph/impl/api.py in wrapper(*args, **kwargs)
      256       except Exception as e:  # pylint:disable=broad-except
      257         if hasattr(e, 'ag_error_metadata'):
  --&gt; 258           raise e.ag_error_metadata.to_exception(e)
      259         else:
      260           raise
  
  ValueError: in user code:
  
      &lt;ipython-input-25-67a2e02c344d&gt;:6 _parse_image  *
          example = tf.io.parse_single_example(example_proto, image_feature, name = &quot;parse_example&quot;)
      /opt/conda/lib/python3.7/site-packages/tensorflow/python/util/dispatch.py:201 wrapper  **
          return target(*args, **kwargs)
      /opt/conda/lib/python3.7/site-packages/tensorflow/python/ops/parsing_ops.py:451 parse_single_example_v2
          serialized = _assert_scalar(serialized, &quot;serialized&quot;)
      /opt/conda/lib/python3.7/site-packages/tensorflow/python/ops/parsing_ops.py:1058 _assert_scalar
          raise ValueError(&quot;Input %s must be a scalar&quot; % name)
  
      ValueError: Input serialized must be a scalar
  </code></pre>
  <p>I'm able to use this TFRecord file in other settings, so I guess my iteration mechanism is faulty.  Please help.</p>