add batch norm

model.py

@@ -13,9 +13,9 @@ img_height = 120
 img_width = 120
 
 # Hyperparameters
-batch_size = 128
-epochs = 10
-learning_rate = 0.001
+batch_size = 64
+epochs = 5
+learning_rate = 0.003
 
 # Create the training dataset
 train_ds = tf.keras.utils.image_dataset_from_directory(
@@ -48,11 +48,13 @@ model = Sequential([
 
     # Encoder
     layers.Conv2D(8, 3, activation='relu'),
+    layers.BatchNormalization(),
     layers.MaxPooling2D(),
     layers.Conv2D(16, 3, activation='relu'),
+    layers.BatchNormalization(),
     layers.MaxPooling2D(),
     layers.Conv2D(32, 3, activation='relu'),
-    # layers.Conv2D(2, 3, activation='relu'), ???
+    layers.BatchNormalization(),
     layers.Flatten(),
 
     # Decoder
@@ -66,7 +68,7 @@ print(model.summary())
 
 # Compile model
 model.compile(loss='categorical_crossentropy',
-              optimizer=keras.optimizers.SGD(learning_rate=learning_rate), metrics=['accuracy'])
+              optimizer=keras.optimizers.Adam(learning_rate=learning_rate), metrics=['accuracy'])
 
 # Train the model on the dataset
 history = model.fit(train_ds, batch_size=batch_size, epochs=epochs,