model_with_batch_norm.py

from tensorflow import keras
from tensorflow.keras import layers

total_pixels = img_size * img_size * 3
translator_factor = 2
translator_layer_size = int(total_pixels/translator_factor)
middle_factor = 2
middle_layer_size = int(translator_layer_size/middle_factor)

inputs = keras.Input(shape=(img_size,img_size,3), name='cat_image')
x = layers.Flatten(name = 'flattened_cat')(inputs) #49152 #3072
x = layers.BatchNormalization()(x)
x = layers.Dense(translator_layer_size, activation='relu', name='encoder')(x)
x = layers.Dense(middle_layer_size, activation='relu', name='middle_layer')(x)
x = layers.BatchNormalization()(x)
x = layers.Dense(translator_layer_size, activation='relu', name='decoder')(x)

outputs = layers.Dense(total_pixels, activation='sigmoid', name='reconstructed_cat')(x)
outputs = layers.Reshape((img_size,img_size,3))(outputs)

model = keras.Model(inputs=inputs, outputs=outputs)