title: Data Augmentation on tf.dataset.Dataset
body: <p>In order to use Google Colabs TPUs I need a <code>tf.dataset.Dataset</code>. How can I then use Data Augmentation on such a dataset? </p>
  
  <p>More specifically, my code so far is:</p>
  
  <pre><code>def get_dataset(batch_size=200):
    datasets, info = tfds.load(name='mnist', with_info=True, as_supervised=True,
                               try_gcs=True)
    mnist_train, mnist_test = datasets['train'], datasets['test']
  
    def scale(image, label):
      image = tf.cast(image, tf.float32)
      image /= 255.0
  
      label = tf.one_hot(label,10)
  
      return image, label
  
    train_dataset = mnist_train.map(scale).shuffle(10000).batch(batch_size)
    test_dataset = mnist_test.map(scale).batch(batch_size)
  
    return train_dataset, test_dataset
  </code></pre>
  
  <p>Which is fed into this:</p>
  
  <pre><code># TPU Strategy ...
  with strategy.scope():
    model = create_model()
    model.compile(loss="categorical_crossentropy",
                  optimizer="adam",
                  metrics=["acc"])
  
  train_dataset, test_dataset = get_dataset()
  
  model.fit(train_dataset,
            epochs=20,
            verbose=1,
            validation_data=test_dataset)
  </code></pre>
  
  <p>So, how can I use here Data Augmentation here?
  As far as I know, I can't use the tf.keras ImageDataGenerator, right?</p>
  
  <p>I've tried the following and it didn't work.</p>
  
  <pre><code>data_generator = ...
  
  model.fit_generator(data_generator.flow(train_dataset, batch_size=32),
                      steps_per_epoch=len(train_dataset) / 32, epochs=20)
  </code></pre>
  
  <p>Not surprising, since, usually, train_x and train_y are fed as two arguments to the flow function, not "packed" into one <code>tf.dataset.Dataset</code>.</p>