Update fomo_utils.py

fomo_utils.py

@@ -220,8 +220,13 @@ def build_fomo_model(input_size, num_classes, alpha=0.35):
     target_size = input_size // 8
     cut_layer = None
     for layer in base.layers:
-        shape = layer.output_shape
-        if isinstance(shape, tuple) and len(shape) == 4 and shape[1] == target_size:
+        try:
+            shape = layer.output.shape
+        except AttributeError:
+            continue
+        if shape is None or len(shape) != 4:
+            continue
+        if shape[1] == target_size:
             cut_layer = layer
     if cut_layer is None:
         raise ValueError(f"Could not find a layer with spatial size {target_size} in MobileNetV2.")
@@ -234,6 +239,15 @@ def build_fomo_model(input_size, num_classes, alpha=0.35):
     return tf.keras.Model(inputs, outputs, name="fomo")
 
 
+def get_optimizer(name, learning_rate):
+    name = (name or "adamw").lower()
+    if name == "adamw":
+        return tf.keras.optimizers.AdamW(learning_rate=learning_rate, weight_decay=1e-4)
+    if name == "sgd":
+        return tf.keras.optimizers.SGD(learning_rate=learning_rate, momentum=0.9)
+    return tf.keras.optimizers.Adam(learning_rate=learning_rate)
+
+
 def weighted_sparse_ce(background_weight):
     def loss_fn(y_true, y_pred):
         y_true = tf.cast(y_true, tf.int32)
@@ -260,9 +274,9 @@ class LogCallback(tf.keras.callbacks.Callback):
         self.log_lines.append(line)
 
 
-def train_model(model, train_ds, val_ds, epochs, learning_rate, background_weight, log_lines):
+def train_model(model, train_ds, val_ds, epochs, learning_rate, background_weight, log_lines, optimizer_name="adamw"):
     model.compile(
-        optimizer=tf.keras.optimizers.Adam(learning_rate),
+        optimizer=get_optimizer(optimizer_name, learning_rate),
         loss=weighted_sparse_ce(background_weight),
         metrics=["accuracy"],
     )