Tensorflow 2.0 Dataloading with aug and caching

data_loading_cell.py

import tensorflow as tf

train_ds = tf.keras.utils.image_dataset_from_directory(
  train_dir,
  validation_split=0.1,
  subset="training",
  label_mode="categorical",
  #shuffle=True,
  seed=123,
  image_size=(image_size, image_size),
  batch_size=batch_size)
class_names = train_ds.class_names

val_ds = tf.keras.utils.image_dataset_from_directory(
  train_dir,
  validation_split=0.1,
  subset="validation",
  label_mode="categorical",
  seed=123,
  image_size=(image_size, image_size),
  batch_size=batch_size)

normalization_layer = tf.keras.layers.Rescaling(1./255)

augmentation = tf.keras.Sequential([
  tf.keras.layers.RandomBrightness(0.2, value_range=(0.0, 1.0)),
  tf.keras.layers.RandomContrast(0.1),
  tf.keras.layers.RandomFlip("horizontal"),
  tf.keras.layers.RandomZoom(0.3)
])


def augmentate(x, y):
  return augmentation(x, training=True), y


AUTOTUNE = tf.data.experimental.AUTOTUNE
train_ds = train_ds.map(lambda x, y: (normalization_layer(x), y)).cache().prefetch(buffer_size=AUTOTUNE).map(augmentate)
val_ds = val_ds.map(lambda x, y: (normalization_layer(x), y)).cache().prefetch(buffer_size=AUTOTUNE)

train_cell_example.py

model.fit_generator(train_ds,
          validation_data=val_ds,
          epochs=100,
          verbose=1)