In Tensorflow, I'm getting shape errors when doing data augmentation with Keras preprocessing layers

Question

I'm training a CNN and am applying data augmentation through Keras layers by defining:

data_augmentation = keras.Sequential([
  layers.experimental.preprocessing.RandomRotation(factor=0.4, fill_mode="wrap"),
  layers.experimental.preprocessing.RandomTranslation(height_factor=0.2, width_factor=0.2, fill_mode="wrap"),
  layers.experimental.preprocessing.RandomFlip("horizontal"),
  layers.experimental.preprocessing.RandomContrast(factor=0.2),
  layers.experimental.preprocessing.RandomHeight(factor=0.2),
  layers.experimental.preprocessing.RandomWidth(factor=0.2)
])

Here are snippets of the code in question:

def process(x, y):
  x = DATASETS_DIR + "/" + x + ".jpg"
  x = tf.io.read_file(x)
  x = tf.image.decode_jpeg(x, channels=3)
  x = tf.image.resize(x, [299, 299])
  x = layers.experimental.preprocessing.Rescaling(1./255)(x)
  
  return x, y

def process_with_augmentation(x, y):
  x, y = process(x, y)

  x = data_augmentation(x)

  return x, y

train_ds = train_ds.map(process_with_augmentation, num_parallel_calls=tf.data.experimental.AUTOTUNE)
validation_ds = validation_ds.map(process, num_parallel_calls=tf.data.experimental.AUTOTUNE)

If I comment out x = data_augmentation(x) in process_with_augmentation() the code works fine. If I don't comment out, I get the following error:

ValueError: Shape must be rank 4 but is rank 3 for '{{node sequential/random_rotation/transform/ImageProjectiveTransformV3}} = ImageProjectiveTransformV3[dtype=DT_FLOAT, fill_mode="WRAP", interpolation="BILINEAR"](rescaling/add, sequential/random_rotation/rotation_matrix/concat, sequential/random_rotation/transform/strided_slice, sequential/random_rotation/transform/fill_value)' with input shapes: [299,299,3], [299,8], [2], [].

Any ideas on how to fix this?

Lescurel Lescurel · Accepted Answer · 2021-01-05T10:06:21

A keras layer expect a batch size. When you map your function on the tf.data.Dataset, the image will lack that batch dimension. You can fix that by adding a dimension before calling your keras preprocessing model :

def process_with_augmentation(x, y):
  x, y = process(x, y)

  x = tf.expand_dims(x, axis=0) # adding the batch dimension
  x = data_augmentation(x)
  x = tf.squeeze(x, [0]) # removing the batch dimension
  return x, y

Another option is to call tf.data.Dataset.batch before calling the model, i.e :

train_ds = train_ds.batch(BATCH_SIZE)
train_ds = train_ds.map(lambda x,y:(data_augmentation(x),y))

If you want to get rid of that extra dimension, you can then call unbatch on the augmented dataset.

This is not the most elegant way of solving this, though. I would suggest integrating directly your preprocessing model in your neural network, or rewriting your preprocessing logic using functions from the tf.image module.

In Tensorflow, I'm getting shape errors when doing data augmentation with Keras preprocessing layers

1 Answers