title: How to apply data augmentation in TensorFlow 2.0 after tfds.load()
body: <p>I'm following <a href="https://www.tensorflow.org/alpha/tutorials/images/transfer_learning#format_the_data" rel="noreferrer">this guide</a>.</p>
  
  <p>It shows how to download datasets from the new TensorFlow Datasets using <code>tfds.load()</code> method:</p>
  
  <pre><code>import tensorflow_datasets as tfds    
  SPLIT_WEIGHTS = (8, 1, 1)
  splits = tfds.Split.TRAIN.subsplit(weighted=SPLIT_WEIGHTS)
  
  (raw_train, raw_validation, raw_test), metadata = tfds.load(
      'cats_vs_dogs', split=list(splits),
      with_info=True, as_supervised=True)
  </code></pre>
  
  <p>The next steps shows how to apply a function to each item in the dataset using map method:</p>
  
  <pre><code>def format_example(image, label):
      image = tf.cast(image, tf.float32)
      image = image / 255.0
      # Resize the image if required
      image = tf.image.resize(image, (IMG_SIZE, IMG_SIZE))
      return image, label
  
  train = raw_train.map(format_example)
  validation = raw_validation.map(format_example)
  test = raw_test.map(format_example)
  </code></pre>
  
  <p>Then to access the elements we can use:</p>
  
  <pre><code>for features in ds_train.take(1):
    image, label = features["image"], features["label"]
  </code></pre>
  
  <p>OR</p>
  
  <pre><code>for example in tfds.as_numpy(train_ds):
    numpy_images, numpy_labels = example["image"], example["label"]
  </code></pre>
  
  <p>However, the guide doesn't mention anything about data augmentation. I want to use real time data augmentation similar to that of Keras's ImageDataGenerator Class. I tried using:</p>
  
  <pre><code>if np.random.rand() &gt; 0.5:
      image = tf.image.flip_left_right(image)
  </code></pre>
  
  <p>and other similar augmentation functions in <code>format_example()</code> but, how can I verify that it's performing real time augmentation and not replacing the original image in the dataset?</p>
  
  <p>I could convert the complete dataset to Numpy array by passing <code>batch_size=-1</code> to <code>tfds.load()</code> and then use <code>tfds.as_numpy()</code> but, that would load all the images in memory which is not needed. I should be able to use <code>train = train.prefetch(tf.data.experimental.AUTOTUNE)</code> to load just enough data for next training loop.</p>