Updated model reconstruction code

app.py

@@ -68,25 +68,17 @@ def create_classification_model():
     # print(inputs)
     
     vision_model.trainable=False
-    # vision_model.trainable = True
-    # # Fine-tune from this layer onwards
-    # fine_tune_at = 100
-
-    # # Freeze all the layers before the `fine_tune_at` layer
-    # for layer in base_model.layers[:fine_tune_at]:
-    # layer.trainable = False
     base_model_output = vision_model(rescaled_input)
     
     current_layer = base_model_output.pooler_output
-    hidden_layers_nodes = [64]
+    hidden_layers_nodes = [1024]
     for node_count in hidden_layers_nodes:
         hidden_layer = tf.keras.layers.Dense(node_count, activation='relu')
         dropout_layer = tf.keras.layers.Dropout(.2, input_shape=(2,))
-        # hidden_layer.trainable = False
         current_layer = hidden_layer(dropout_layer(current_layer))
     
-    prediction_layer = tf.keras.layers.Dense(9, activation='softmax')
-    # prediction_layer.trainable = False
+    prediction_layer = tf.keras.layers.Dense(
+        classes_count, activation='softmax')
     outputs = prediction_layer(current_layer)
     model = tf.keras.Model(inputs, outputs)
 
@@ -94,10 +86,11 @@ def create_classification_model():
     model.compile(
         # Used leagcy optimizer due to tf 2.11 release issues with MACOS 
         # optimizer=tf.keras.optimizers.Adam(learning_rate=base_learning_rate),
-        optimizer=tf.keras.optimizers.legacy.Adam(learning_rate=base_learning_rate),
+        optimizer=tf.keras.optimizers.legacy.Adam(
+            learning_rate=base_learning_rate),
         loss=tf.keras.losses.SparseCategoricalCrossentropy(),
-        metrics=['accuracy'],
-        steps_per_execution=steps_per_execution
+        metrics=['accuracy']
+        # steps_per_execution=steps_per_execution
     )
 
     return model