title: Why would this dataset implementation run out of memory?
body: <p>I follow <a href="https://www.tensorflow.org/tutorials/load_data/images" rel="noreferrer">this instruction</a> and write the following code to create a Dataset for images(<a href="http://msvocds.blob.core.windows.net/coco2014/train2014.zip" rel="noreferrer">COCO2014 training set</a>)</p>
  
  <pre class="lang-py prettyprint-override"><code>from pathlib import Path
  import tensorflow as tf
  
  
  def image_dataset(filepath, image_size, batch_size, norm=True):
      def preprocess_image(image):
          image = tf.image.decode_jpeg(image, channels=3)
          image = tf.image.resize(image, image_size)
          if norm:
              image /= 255.0  # normalize to [0,1] range
          return image
  
      def load_and_preprocess_image(path):
          image = tf.read_file(path)
          return preprocess_image(image)
  
      all_image_paths = [str(f) for f in Path(filepath).glob('*')]
      path_ds = tf.data.Dataset.from_tensor_slices(all_image_paths)
      ds = path_ds.map(load_and_preprocess_image, num_parallel_calls=tf.data.experimental.AUTOTUNE)
      ds = ds.shuffle(buffer_size = len(all_image_paths))
      ds = ds.repeat()
      ds = ds.batch(batch_size)
      ds = ds.prefetch(tf.data.experimental.AUTOTUNE)
  
      return ds
  
  ds = image_dataset(train2014_dir, (256, 256), 4, False)
  image = ds.make_one_shot_iterator().get_next('images')
  # image is then fed to the network
  </code></pre>
  
  <p>This code will always run out of both memory(32G) and GPU(11G) and kill the process. Here is the messages shown on terminal.
  <a href="https://i.stack.imgur.com/FQ3Df.png" rel="noreferrer"><img src="https://i.stack.imgur.com/FQ3Df.png" alt="enter image description here"></a></p>
  
  <p>I also spot that the program get stuck at <code>sess.run(opt_op)</code>. Where is wrong? How can I fix it?</p>