Implement Two-Stage Fine-Tuning

Now pipeline 2 precisely follows ValizadehAslani et al.'s Two-Stage Fine-Tuning.

Author: lucasalvaa

pipeline2/dvc.yaml

@@ -1,23 +1,11 @@
 stages:
-  balancing:
-    wdir: ..
-    cmd: python -m src.preprocessing.balance --config pipeline2/params.yaml
-    deps:
-      - src/preprocessing/balance.py
-      - data/split
-    params:
-      - pipeline2/params.yaml:
-        - data
-    outs:
-      - data_balanced
-
-  train:
+  stage-1:
     foreach: [ effnet_s, effnet_m, convnext ]
     do:
       wdir: ..
-      cmd: python -m src.train --model ${item} --config pipeline2/params.yaml --output_dir pipeline2/${item}
+      cmd: python -m src.train --model ${item} --config pipeline2/params.yaml --output_dir pipeline2/${item} --tsft true
       deps:
-        - data_balanced
+        - data/split/train
         - data/split/val
         - src/train.py
         - src/common.py
@@ -30,41 +18,16 @@ stages:
         - pipeline2/${item}/model.pth
         - pipeline2/${item}/loss.json
 
-  evaluate:
-      foreach: [ effnet_s, effnet_m, convnext ]
-      do:
-        wdir: ..
-        cmd: python -m src.evaluate --model ${item} --config pipeline2/params.yaml --output_dir pipeline2/${item}
-        deps:
-          - data/split/test
-          - src/evaluate.py
-          - src/common.py
-          - pipeline2/${item}/model.pth
-        params:
-          - pipeline2/params.yaml:
-            - base
-            - data
-            - evaluate
-        metrics:
-          - pipeline2/${item}/metrics.json:
-              cache: false
-        plots:
-          - pipeline2/${item}/cm_data.csv:
-              template: confusion
-              x: actual
-              y: predicted
-              title: "Pipeline 2 - Balance CM - ${item}"
-              cache: false
-
-  finetuning:
+  stage-2: # fine-tuning
     foreach: [ effnet_s, effnet_m, convnext ]
     do:
       wdir: ..
       cmd: python -m src.finetune --model ${item} --config pipeline2/params.yaml --output_dir pipeline2/${item}/finetuned
       deps:
+        - data/split/train
+        - data/split/val
         - src/finetune.py
         - src/common.py
-        - data/split/train
         - pipeline2/${item}/model.pth
       params:
         - pipeline2/params.yaml:
@@ -76,7 +39,7 @@ stages:
         - pipeline2/${item}/finetuned/loss.json
 
 
-  ft-evaluate:
+  evaluate:
     foreach: [ effnet_s, effnet_m, convnext ]
     do:
       wdir: ..
@@ -99,21 +62,11 @@ stages:
             template: confusion
             x: actual
             y: predicted
-            title: "Pipeline 2 - Balance + Finetune CM - ${item}"
+            title: "Pipeline 2 - Two-Stage Fine-Tuning CM - ${item}"
             cache: false
 
 
 plots:
-    - Training_Loss_Comparison:
-        template: linear
-        x: epoch
-        y:
-          # Qui confrontiamo le performance di addestramento tra i modelli
-          effnet_s/loss.json: train_loss
-          effnet_m/loss.json: train_loss
-          convnext/loss.json: train_loss
-        title: "Pipeline 2 - Comparison: Training Loss per Model"
-
     - effnet_s_curves:
         template: linear
         x: epoch

pipeline2/params.yaml

@@ -2,9 +2,7 @@ base:
   image_res: 224
 
 data:
-  tobalance_path: data/split
-  balanced_path: data_balanced/train
-  trainset_path: data_balanced/train
+  trainset_path: data/split/train
   valset_path: data/split/val
   testset_path: data/split/test
 

src/common.py

@@ -1,9 +1,11 @@
-"""Common utilities for the P1 pipeline, including data loading and model setup."""
+"""Common utilities."""
 
 from pathlib import Path
 
 import torch
+import torch.optim as optim
 from torch import nn
+from torch.amp import GradScaler, autocast
 from torch.utils.data import DataLoader
 from torchvision import datasets, models, transforms
 
@@ -37,7 +39,7 @@ def get_dataloader(
         dataset,
         batch_size=batch_size,
         shuffle=("train" in str(data_path)),
-        num_workers=4,  # Consiglio: accelera il caricamento dati
+        num_workers=4,  # Accelera il caricamento dati
         pin_memory=True,  # Accelera il trasferimento dati alla GPU
     )
 
@@ -65,3 +67,64 @@ def get_model(model_name: str, num_classes: int) -> nn.Module:
         model.classifier[2] = nn.Linear(model.classifier[2].in_features, num_classes)
 
     return model.to(DEVICE)
+
+def validate(model: nn.Module, loader: DataLoader, criterion: nn.Module) -> float:
+    """Calculate average loss on the validation set.
+
+    Args:
+        model: The neural network model.
+        loader: DataLoader for the validation set.
+        criterion: Loss function (e.g., CrossEntropyLoss or LDAMLoss).
+
+    Returns:
+        The average loss over the entire dataset.
+
+    """
+    model.eval()
+    running_loss = 0.0
+    with torch.no_grad():
+        for images, targets in loader:
+            images, targets = images.to(DEVICE), targets.to(DEVICE)
+            outputs = model(images)
+            loss = criterion(outputs, targets)
+            running_loss += loss.item() * images.size(0)
+    return running_loss / len(loader.dataset)
+
+def train_epoch(
+    model: nn.Module,
+    loader: DataLoader,
+    criterion: nn.Module,
+    optimizer: optim.Optimizer,
+    scaler: GradScaler,
+) -> float:
+    """Run one training epoch with AMP.
+
+    Args:
+        model: The neural network model.
+        loader: DataLoader for the training set.
+        criterion: Loss function.
+        optimizer: Optimizer.
+        scaler: GradScaler for AMP.
+
+    Returns:
+        Average training loss.
+
+    """
+    model.train()
+    running_loss = 0.0
+
+    for images, targets in loader:
+        images, targets = images.to(DEVICE), targets.to(DEVICE)
+        optimizer.zero_grad()
+
+        with autocast(device_type=DEVICE.type):
+            outputs = model(images)
+            loss = criterion(outputs, targets)
+
+        scaler.scale(loss).backward()
+        scaler.step(optimizer)
+        scaler.update()
+
+        running_loss += loss.item() * images.size(0)
+
+    return running_loss / len(loader.dataset)

src/finetune.py

@@ -8,104 +8,8 @@
 import torch
 import yaml
 from torch import amp, nn, optim
-from torch.utils.data import DataLoader
 
-# Assumendo che la struttura dei package sia corretta rispetto alla root
-from src.common import DEVICE, get_dataloader, get_model
-
-
-def find_best_model(models_list: list[str], base_path: Path = Path(".")) -> str:
-    """Identify the model with the highest top1 metric.
-
-    Args:
-        models_list: List of model names to check.
-        base_path: Directory containing model folders.
-
-    Returns:
-        The name of the best model.
-
-    """
-    best_top1 = -1.0
-    best_model_name = ""
-
-    for model_name in models_list:
-        metrics_path = base_path / model_name / "metrics.json"
-        if metrics_path.exists():
-            with metrics_path.open("r") as f:
-                data = json.load(f)
-                if data["top1"] > best_top1:
-                    best_top1 = data["top1"]
-                    best_model_name = model_name
-
-    if not best_model_name:
-        raise FileNotFoundError(
-            "No valid metrics.json found to determine the best model."
-        )
-
-    return best_model_name
-
-
-def validate(model: nn.Module, loader: DataLoader, criterion: nn.Module) -> float:
-    """Calculate average loss on the validation set.
-
-    Args:
-        model: The neural network model.
-        loader: DataLoader for validation.
-        criterion: Loss function.
-
-    Returns:
-        Average validation loss.
-
-    """
-    model.eval()
-    running_loss = 0.0
-    with torch.no_grad():
-        for images, targets in loader:
-            images, targets = images.to(DEVICE), targets.to(DEVICE)
-            outputs = model(images)
-            loss = criterion(outputs, targets)
-            running_loss += loss.item() * images.size(0)
-    return running_loss / len(loader.dataset)
-
-
-def train_one_epoch(
-    model: nn.Module,
-    loader: DataLoader,
-    criterion: nn.Module,
-    optimizer: optim.Optimizer,
-    scaler: amp.GradScaler,
-) -> float:
-    """Run one fine-tuning epoch using AMP.
-
-    Args:
-        model: The neural network model.
-        loader: DataLoader for the training set.
-        criterion: Loss function.
-        optimizer: Optimizer.
-        scaler: GradScaler for AMP.
-
-    Returns:
-        Average training loss.
-
-    """
-    model.train()
-    running_loss = 0.0
-
-    for images, targets in loader:
-        images, targets = images.to(DEVICE), targets.to(DEVICE)
-        optimizer.zero_grad()
-
-        with amp.autocast(device_type="cuda" if torch.cuda.is_available() else "cpu"):
-            outputs = model(images)
-            loss = criterion(outputs, targets)
-
-        scaler.scale(loss).backward()
-        scaler.step(optimizer)
-        scaler.update()
-
-        running_loss += loss.item() * images.size(0)
-
-    return running_loss / len(loader.dataset)
+from src.common import DEVICE, get_dataloader, get_model, train_epoch, validate
 
 
 def main() -> None:
@@ -120,54 +24,55 @@ def main() -> None:
     with open(args.config) as conf_file:
         config: dict[str, Any] = yaml.safe_load(conf_file)
 
-    # 2. Setup directory di output (sovrascriviamo o creiamo una cartella fine_tuned)
     out_dir = Path(args.output_dir)
     out_dir.mkdir(parents=True, exist_ok=True)
 
-    # 3. Caricamento dati (Phase 2 per fine-tuning)
-    # Nota: Assumiamo che phase2 sia la directory dei dati di training bilanciati
-    train_loader = get_dataloader(
+
+    t_loader = get_dataloader(
         data_path=Path(config["finetuning"]["data_path"]),
         batch_size=config["finetuning"]["batch_size"],
     )
 
-    # Usiamo il set di validazione originale per il monitoraggio
-    val_loader = get_dataloader(
+    v_loader = get_dataloader(
         data_path=Path(config["data"]["valset_path"]),
         batch_size=config["finetuning"]["batch_size"],
     )
 
-    # 4. Inizializzazione modello e caricamento pesi precedenti
-    model = get_model(args.model, len(train_loader.dataset.classes))
+    # Model initialization loading first stage's weights
+    model = get_model(args.model, len(t_loader.dataset.classes))
     weights_path = out_dir.parent / "model.pth"
-    model.load_state_dict(torch.load(weights_path, map_location=DEVICE))
-    model.to(DEVICE)
+    model.load_state_dict(torch.load(weights_path, map_location=DEVICE)).to(DEVICE)
 
-    # 5. Configurazione training
-    # Per il fine-tuning si usa solitamente un Learning Rate più basso (es. 1e-5 o 1e-4)
+    # Unfreeze layers
+    for param in model.parameters():
+        param.requires_grad = True
 
+    # Fine-tuning setup
     criterion = nn.CrossEntropyLoss()
     optimizer = optim.Adam(model.parameters(), lr=config["finetuning"]["lr"])
     scaler = amp.GradScaler()
 
+    # Model fine-tuning
     history = []
-
-    # 6. Loop di fine-tuning
     epochs = config["finetuning"]["epochs"]
+    print(f"Fine-tuning {args.model}...")
     for epoch in range(epochs):
-        t_loss = train_one_epoch(model, train_loader, criterion, optimizer, scaler)
-        v_loss = validate(model, val_loader, criterion)
+        t_loss = train_epoch(model, t_loader, criterion, optimizer, scaler)
+        v_loss = validate(model, v_loader, criterion)
 
         history.append({"epoch": epoch + 1, "train_loss": t_loss, "val_loss": v_loss})
 
         print(
             f"Epoch {epoch + 1}/{epochs} | "
-            f"FT Train Loss: {t_loss:.4f} | "
-            f"FT Val Loss: {v_loss:.4f}"
+            f"T-Loss: {t_loss:.4f} | "
+            f"V-Loss: {v_loss:.4f}"
         )
+    print(f"Model {args.model} fine-tuned successfully!")
 
-    # 7. Salvataggio artefatti
+    # Saving the model
     torch.save(model.state_dict(), out_dir / "model.pth")
+
+    # Saving training and validation loss in loss.json file
     with open(out_dir / "loss.json", "w") as f:
         json.dump(history, f, indent=4)