Data Augmentation on tf.dataset.Dataset

Question

In order to use Google Colabs TPUs I need a tf.dataset.Dataset. How can I then use Data Augmentation on such a dataset?

More specifically, my code so far is:

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

Which is fed into this:

# 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)

So, how can I use here Data Augmentation here? As far as I know, I can't use the tf.keras ImageDataGenerator, right?

I've tried the following and it didn't work.

data_generator = ...

model.fit_generator(data_generator.flow(train_dataset, batch_size=32),
                    steps_per_epoch=len(train_dataset) / 32, epochs=20)

Not surprising, since, usually, train_x and train_y are fed as two arguments to the flow function, not "packed" into one tf.dataset.Dataset.

user_13465258 · Answer

You can use tf.image functions. The tf.image module contains various functions for image processing.

For example:

You can add below functionality in your function def get_dataset.

convert each image to tf.float64 in the 0-1 range.
cache() results as those can be re-used after each repeat
randomly flip left_to_right each image using random_flip_left_right.
randomly change contrast of image using random_contrast.
Number of images increased by twice by repeat which repeat all the steps.

Code -

mnist_train = mnist_train.map(
    lambda image, label: (tf.image.convert_image_dtype(image, tf.float32), label)
).cache(
).map(
    lambda image, label: (tf.image.random_flip_left_right(image), label)
).map(
    lambda image, label: (tf.image.random_contrast(image, lower=0.0, upper=1.0), label)
).shuffle(
    1000
).
batch(
    batch_size
).repeat(2)

Similarly you can use other functionalities like random_flip_up_down, random_crop functions to Randomly flips an image vertically (upside down) and Randomly crop a tensor to a given size respectively.

Your get_dataset function will look like below -

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']

  train_dataset = mnist_train.map(
               lambda image, label: (tf.image.convert_image_dtype(image, tf.float32),label)
              ).cache(
              ).map(
                    lambda image, label: (tf.image.random_flip_left_right(image), label)
              ).map(
                    lambda image, label: (tf.image.random_contrast(image, lower=0.0, upper=1.0), label)
              ).shuffle(
                    1000
              ).batch(
                    batch_size
              ).repeat(2)

  test_dataset = mnist_test.map(scale).batch(batch_size)

  return train_dataset, test_dataset

Adding the link suggested by @Andrew H that gives end-to-end example on data augmentation that also uses mnist dataset.

Hope this answers your question. Happy Learning.

Data Augmentation on tf.dataset.Dataset

1 Answers