Updated Model Architecture, Removed EfficientNet Support

Author: animikhaich

main.py

@@ -14,6 +14,7 @@
 # TODO: Add Support For Learning Rate Change
 # TODO: Add Support For Dynamic Polt.ly Charts
 # TODO: Add Support For Live Training Graphs (on_train_batch_end) without slowing down the Training Process
+# TODO: Add Supoort For EfficientNet - Fix Data Loader Input to be Un-Normalized Images
 
 OPTIMIZERS = {
     "SGD": tf.keras.optimizers.SGD(),
@@ -28,6 +29,27 @@
 
 BATCH_SIZES = [1, 2, 4, 8, 16, 32, 64, 128, 256]
 
+BACKBONES = [
+    "MobileNet",
+    "MobileNetV2",
+    "ResNet50",
+    "ResNet101",
+    "ResNet152",
+    "ResNet50V2",
+    "ResNet101V2",
+    "ResNet152V2",
+    "VGG16",
+    "VGG19",
+    "Xception",
+    "InceptionV3",
+    "InceptionResNetV2",
+    "DenseNet121",
+    "DenseNet169",
+    "DenseNet201",
+    "NASNetMobile",
+    "NASNetLarge",
+]
+
 
 class CustomCallback(tf.keras.callbacks.Callback):
     def __init__(self, num_steps):
@@ -69,11 +91,13 @@ def on_epoch_begin(self, epoch, logs=None):
 
     def on_train_begin(self, logs=None):
         self.status_text.info(
-            "Training Started! Live Graphs will be shown on the completion of the First Epoch"
+            "Training Started! Live Graphs will be shown on the completion of the First Epoch."
         )
 
     def on_train_end(self, logs=None):
-        self.status_text.success("Training Completed!")
+        self.status_text.success(
+            f"Training Completed! Final Validation Accuracy: {logs['val_categorical_accuracy']*100:.2f}%"
+        )
         st.balloons()
 
     def on_epoch_end(self, epoch, logs=None):
@@ -126,14 +150,20 @@ def on_epoch_end(self, epoch, logs=None):
         "Tensorboard Logs Directory (Absolute Path)", "logs/tensorboard"
     )
 
+    # Select Backbone
+    selected_backbone = st.selectbox("Select Backbone", BACKBONES)
+
     # Select Optimizer
     selected_optimizer = st.selectbox("Training Optimizer", list(OPTIMIZERS.keys()))
 
     # Select Batch Size
     selected_batch_size = st.select_slider("Train/Eval Batch Size", BATCH_SIZES, 16)
 
     # Select Number of Epochs
-    selected_epochs = st.number_input("Max Number of Epochs", 1, 300000, 100)
+    selected_epochs = st.number_input("Max Number of Epochs", 1, 500, 100)
+
+    # Select Number of Epochs
+    selected_input_shape = st.number_input("Input Image Shape", 64, 2000, 224)
 
     # Start Training Button
     start_training = st.button("Start Training")
@@ -161,7 +191,7 @@ def on_epoch_end(self, epoch, logs=None):
     )
 
     classifier = ImageClassifier(
-        backbone="EfficientNetB0",
+        backbone=selected_backbone,
         input_shape=(224, 224, 3),
         classes=train_data_loader.get_num_classes(),
     )
@@ -170,8 +200,6 @@ def on_epoch_end(self, epoch, logs=None):
     classifier.set_tensorboard_path(tensorboard_logs_path)
 
     classifier.init_callbacks(
-        keras_weights_path,
-        tensorboard_logs_path,
         [CustomCallback(train_data_loader.get_num_steps())],
     )
 

utils/model.py

@@ -99,8 +99,10 @@ def get_default_callbacks(self, weights_path, tb_logs_path):
         callbacks = [early_stop_cb, model_ckpt_cb, reduce_lr_cb, tensorboard_cb]
         return callbacks
 
-    def init_callbacks(self, weights_path, tb_logs_path, custom_callbacks=[]):
-        self.callbacks = self.get_default_callbacks(weights_path, tb_logs_path)
+    def init_callbacks(self, custom_callbacks=[]):
+        self.callbacks = self.get_default_callbacks(
+            self.keras_weights_path, self.tensorboard_logs_path
+        )
         self.callbacks.extend(custom_callbacks)
         return self.callbacks
 
@@ -180,8 +182,6 @@ def init_network(self):
             backbone=self.backbone, input_shape=self.input_shape, classes=self.classes
         )
         x = base.output
-        x = tf.keras.layers.Conv2D(64, (1, 1), activation="relu")(x)
-        x = tf.keras.layers.Dropout(0.5)(x)
         x = tf.keras.layers.Flatten()(x)
         x = tf.keras.layers.Dense(units=128, activation="relu")(x)
         x = tf.keras.layers.Dropout(0.5)(x)