Input of the model has been standardized to BWHF (Batch, Width, Height, Features) shape. RandomScale2DLayer and RandomShift2DLayer has been incorporated into the model in order to apply these two steps of data augmentation on-the-fly during the training of the model, but not during the test of the model.