Update train_utils.py

train_utils.py

@@ -55,6 +55,7 @@ def load_model(model_name: str, device: torch.device) -> Tuple[nn.Module, dict]:
 
     if not os.path.exists(meta_file):
         raise FileNotFoundError(f"Métadonnées introuvables pour le modèle : {model_name}")
+
     if not os.path.exists(weight_file):
         raise FileNotFoundError(f"Poids introuvables pour le modèle : {model_name}")
 
@@ -65,9 +66,9 @@ def load_model(model_name: str, device: torch.device) -> Tuple[nn.Module, dict]:
 
     model = ResNet18Classifier(
         num_classes=cfg["num_classes"],
-        dropout=cfg["dropout"],
-        fc_dim=cfg["fc_dim"],
-        freeze_backbone=cfg.get("freeze_backbone", True),
+        dropout=cfg.get("dropout", 0.4),
+        fc_dim=cfg.get("fc_dim", 256),
+        fine_tune_mode=cfg.get("fine_tune_mode", "layer4"),
     )
 
     state_dict = torch.load(weight_file, map_location="cpu")
@@ -113,6 +114,7 @@ def collect_predictions(model, loader, device):
     with torch.no_grad():
         for images, labels in loader:
             images = images.to(device)
+
             outputs = model(images)
             preds = outputs.argmax(dim=1).detach().cpu().tolist()
 
@@ -123,14 +125,14 @@ def collect_predictions(model, loader, device):
 
 
 def train_model(
-    dropout: float,
-    fc_dim: int,
-    learning_rate: float,
-    weight_decay: float,
-    batch_size: int,
-    epochs: int,
-    freeze_backbone: bool,
-    model_tag: str,
+    dropout: float = 0.4,
+    fc_dim: int = 256,
+    learning_rate: float = 0.00001,
+    weight_decay: float = 0.0001,
+    batch_size: int = 16,
+    epochs: int = 30,
+    fine_tune_mode: str = "layer4",
+    model_tag: str = "",
 ):
     device = get_runtime_device()
 
@@ -141,13 +143,14 @@ def train_model(
         num_classes=num_classes,
         dropout=dropout,
         fc_dim=fc_dim,
-        freeze_backbone=freeze_backbone,
+        fine_tune_mode=fine_tune_mode,
     ).to(device)
 
     trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
     total_params = sum(p.numel() for p in model.parameters())
 
     criterion = nn.CrossEntropyLoss()
+
     optimizer = optim.AdamW(
         filter(lambda p: p.requires_grad, model.parameters()),
         lr=learning_rate,
@@ -176,6 +179,10 @@ def train_model(
 
             loss = criterion(outputs, labels)
             loss.backward()
+
+            # Important: prevents unstable fine-tuning / exploding gradients
+            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+
             optimizer.step()
 
             running_loss += loss.item() * images.size(0)
@@ -239,13 +246,14 @@ def train_model(
         "weight_decay": weight_decay,
         "batch_size": batch_size,
         "epochs": epochs,
-        "freeze_backbone": freeze_backbone,
+        "fine_tune_mode": fine_tune_mode,
     }
 
     training_summary = {
         "final_train_loss": history[-1]["train_loss"] if history else None,
         "final_train_acc": history[-1]["train_acc"] if history else None,
         "best_val_loss": round(best_val_loss, 4),
+        "final_val_loss": history[-1]["val_loss"] if history else None,
         "final_val_acc": history[-1]["val_acc"] if history else None,
         "test_cross_entropy_loss": round(test_loss, 4),
         "test_accuracy": round(test_acc, 4),
@@ -263,6 +271,7 @@ def train_model(
     logs.append("Entraînement terminé.")
     logs.append(f"Modèle sauvegardé : {model_name}")
     logs.append(f"Appareil utilisé : {device}")
+    logs.append(f"Mode de fine-tuning : {fine_tune_mode}")
     logs.append(f"Nombre total de paramètres : {total_params}")
     logs.append(f"Paramètres entraînables : {trainable_params}")
     logs.append(f"Perte test cross-entropy : {test_loss:.4f}")
@@ -289,7 +298,7 @@ def evaluate_saved_model(model_name: str):
     device = get_runtime_device()
     model, meta = load_model(model_name, device)
 
-    batch_size = int(meta["config"].get("batch_size", 32))
+    batch_size = int(meta["config"].get("batch_size", 16))
     _, _, test_loader, class_names = make_loaders(batch_size)
 
     criterion = nn.CrossEntropyLoss()