Add early stopping and minor changes in dvc.yaml files.

Co-authored-by: SimoCimmi <simonecimmino2004@gmail.com>
Co-authored-by: Morgan Vitiello <morgan.vitiello06@gmail.com>

Author: 3 people

baseline/dvc.yaml

@@ -3,7 +3,7 @@ stages:
     foreach: [ effnet_s, effnet_m, convnext ]
     do:
       wdir: ..
-      cmd: python -m src.train --model ${item} --config baseline/params.yaml --output_dir baseline/${item}
+      cmd: python -m src.train --pipeline baseline --model ${item}
       deps:
         - data/split/train
         - data/split/val
@@ -22,7 +22,7 @@ stages:
     foreach: [ effnet_s, effnet_m, convnext ]
     do:
       wdir: ..
-      cmd: python -m src.evaluate --model ${item} --config baseline/params.yaml --output_dir baseline/${item}
+      cmd: python -m src.evaluate --pipeline baseline --model ${item}
       deps:
         - data/split/test
         - src/evaluate.py

pipeline1/dvc.yaml

@@ -15,7 +15,7 @@ stages:
     foreach: [ effnet_s, effnet_m, convnext ]
     do:
       wdir: ..
-      cmd: python -m src.train --model ${item} --config pipeline1/params.yaml --output_dir pipeline1/${item}
+      cmd: python -m src.train --pipeline pipeline1 --model ${item}
       deps:
         - data_augmented/train
         - data_augmented/val
@@ -34,7 +34,7 @@ stages:
     foreach: [ effnet_s, effnet_m, convnext ]
     do:
       wdir: ..
-      cmd: python -m src.evaluate --model ${item} --config pipeline1/params.yaml --output_dir pipeline1/${item}
+      cmd: python -m src.evaluate --pipeline pipeline1 --model ${item}
       deps:
         - data_augmented/test
         - src/evaluate.py

pipeline2/dvc.yaml

@@ -3,7 +3,7 @@ stages:
     foreach: [ effnet_s, effnet_m, convnext ]
     do:
       wdir: ..
-      cmd: python -m src.train --model ${item} --config pipeline2/params.yaml --output_dir pipeline2/${item} --tsft true
+      cmd: python -m src.train --pipeline pipeline2 --model ${item} --tsft true
       deps:
         - data/split/train
         - data/split/val
@@ -21,7 +21,7 @@ stages:
     foreach: [ effnet_s, effnet_m, convnext ]
     do:
       wdir: ..
-      cmd: python -m src.finetune --model ${item} --config pipeline2/params.yaml --output_dir pipeline2/${item}/finetuned
+      cmd: python -m src.finetune --pipeline pipeline2 --model ${item}
       deps:
         - data/split/train
         - data/split/val
@@ -42,7 +42,7 @@ stages:
     foreach: [ effnet_s, effnet_m, convnext ]
     do:
       wdir: ..
-      cmd: python -m src.evaluate --model ${item} --config pipeline2/params.yaml --output_dir pipeline2/${item}/finetuned
+      cmd: python -m src.evaluate --pipeline pipeline2 --model ${item}
       deps:
         - data/split/test
         - src/evaluate.py

pipeline3/dvc.yaml

@@ -1,82 +1,44 @@
 stages:
-  balancing:
-    wdir: ..
-    cmd: python -m src.preprocessing.balance --config pipeline3/params.yaml
-    deps:
-      - src/preprocessing/balance.py
-      - data/split
-    params:
-      - pipeline3/params.yaml:
-        - data
-    outs:
-      - pipeline3/data_balanced
-
-  augment:
+  augmentation:
     wdir: ..
     cmd: python -m src.preprocessing.augment --config pipeline3/params.yaml
     deps:
+      - data/split
       - src/preprocessing/balance.py
-      - pipeline3/data_balanced
     params:
       - pipeline3/params.yaml:
           - data
-    outs:
-      - pipeline3/data_augmented
+    # outs:
+      # - pipeline3/data_augmented
 
-  train:
+  stage-1:
     foreach: [ effnet_s, effnet_m, convnext ]
     do:
       wdir: ..
-      cmd: python -m src.train --model ${item} --config pipeline3/params.yaml --output_dir pipeline3/${item}
+      cmd: python -m src.train --pipeline pipeline3 --model ${item} --tsft true
       deps:
-        - pipeline3/data_augmented
+        - data_augmented/train
         - data/split/val
         - src/train.py
         - src/common.py
       params:
         - pipeline3/params.yaml:
-          - base
-          - data
-          - train
-      outs:
-        - pipeline3/${item}/model.pth
-        - pipeline3/${item}/loss.json
-
-  evaluate:
-      foreach: [ effnet_s, effnet_m, convnext ]
-      do:
-        wdir: ..
-        cmd: python -m src.evaluate --model ${item} --config pipeline3/params.yaml --output_dir pipeline3/${item}
-        deps:
-          - data/split/test
-          - src/evaluate.py
-          - src/common.py
-          - pipeline3/${item}/model.pth
-        params:
-          - pipeline3/params.yaml:
             - base
             - data
-            - evaluate
-        metrics:
-          - pipeline3/${item}/metrics.json:
-              cache: false
-        plots:
-          - pipeline3/${item}/cm_data.csv:
-              template: confusion
-              x: actual
-              y: predicted
-              title: "Pipeline 3 - Balance + Augment CM - ${item}"
-              cache: false
+            - train
+      outs:
+        - pipeline3/${item}/model.pth
 
-  finetuning:
+  stage-2: # fine-tuning
     foreach: [ effnet_s, effnet_m, convnext ]
     do:
       wdir: ..
-      cmd: python -m src.finetune --model ${item} --config pipeline3/params.yaml --output_dir pipeline3/${item}/finetuned
+      cmd: python -m src.finetune --pipeline pipeline3 --model ${item}
       deps:
+        - pipeline3/data_augmented/train
+        - data/split/val
         - src/finetune.py
         - src/common.py
-        - data/split/train
         - pipeline3/${item}/model.pth
       params:
         - pipeline3/params.yaml:
@@ -87,34 +49,32 @@ stages:
         - pipeline3/${item}/finetuned/model.pth
         - pipeline3/${item}/finetuned/loss.json
 
-
-  ft-evaluate:
+  evaluate:
     foreach: [ effnet_s, effnet_m, convnext ]
     do:
       wdir: ..
-      cmd: python -m src.evaluate --model ${item} --config pipeline3/params.yaml --output_dir pipeline3/${item}/finetuned
+      cmd: python -m src.evaluate --pipeline pipeline3 --model ${item}
       deps:
         - data/split/test
         - src/evaluate.py
         - src/common.py
         - pipeline3/${item}/finetuned/model.pth
       params:
         - pipeline3/params.yaml:
-          - base
-          - data
-          - evaluate
+            - base
+            - data
+            - evaluate
       metrics:
         - pipeline3/${item}/finetuned/metrics.json:
             cache: false
       plots:
-        - pipeline3/${item}/finetuned/cm_data.json:
+        - pipeline3/${item}/finetuned/cm_data.csv:
             template: confusion
             x: actual
             y: predicted
-            title: "Pipeline 3 - Balance + Augment + Finetune CM - ${item}"
+            title: "Pipeline 3 - Two-Stage Fine-Tuning CM - ${item}"
             cache: false
 
-
 plots:
     - Training_Loss_Comparison:
         template: linear

src/early_stopping.py

@@ -0,0 +1,45 @@
+import torch
+import torch.nn as nn
+
+
+class EarlyStopping:
+    """Implement the Automatic Early Stopping technique (Lutz Prechelt, 1998)
+    In particular, it uses the GL_alpha criterion: the training stops
+    when the Generalization Loss is greater than the alpha value.
+    """
+
+    def __init__(self, alpha: float = 5.0, path: str = "checkpoint.pth") -> None:
+        """Inizializza il monitoraggio.
+
+        Args:
+            alpha: Soglia percentuale di Generalization Loss (es. 5.0).
+            path: Percorso dove salvare il miglior modello (E_opt).
+
+        """
+        self.alpha: float = alpha
+        self.path: str = path
+        self.min_v_loss: float = float("inf")
+        self.best_epoch: int = 0
+        self.stop: bool = False
+
+    def __call__(self, v_loss: float, epoch: int, model: nn.Module) -> None:
+        """Verifica la condizione di arresto.
+
+        Args:
+            v_loss: Loss di validazione dell'epoca corrente.
+            epoch: Indice dell'epoca attuale.
+            model: Il modello da salvare in caso di miglioramento.
+
+        """
+        if v_loss < self.min_v_loss:
+            self.min_v_loss = v_loss
+            self.best_epoch = epoch
+            # Salviamo il modello "ottimale" (E_opt) citato nel paper
+            torch.save(model.state_dict(), self.path)
+
+        # GL(t) = 100 * (E_va(t) / E_opt(t) - 1)
+        gl_t = 100 * (v_loss / self.min_v_loss - 1)
+
+        if gl_t > self.alpha:
+            print(f"\n[Early Stopping] GL: {gl_t:.2f}% > Alpha: {self.alpha}%")
+            self.stop = True

src/evaluate.py

@@ -53,18 +53,18 @@ def evaluate(
 
 def main() -> None:
     """Run test evaluation and save artifacts."""
+    choices = ["baseline", "pipeline1", "pipeline2", "pipeline3"]
     parser = argparse.ArgumentParser()
+    parser.add_argument("--pipeline", choices=choices, type=str, required=True)
     parser.add_argument("--model", type=str, required=True)
-    parser.add_argument("--config", type=str, required=True)
-    parser.add_argument("--output_dir", type=str, required=True)
     parser.add_argument("--model_path", type=str, default=None)
-
     args = parser.parse_args()
 
-    with open(args.config) as conf_file:
-        config = yaml.safe_load(conf_file)
+    params_path = Path(args.pipeline) / "params.yaml"
+    with open(params_path) as f:
+        config = yaml.safe_load(f)
 
-    out_dir = Path(args.output_dir)
+    out_dir = Path(args.pipeline) / Path(args.model)
     out_dir.mkdir(parents=True, exist_ok=True)
 
     test_loader = get_dataloader(

src/finetune.py

@@ -3,28 +3,30 @@
 import argparse
 import json
 from pathlib import Path
-from typing import Any
 
 import torch
 import yaml
 from torch import amp, nn, optim
 
 from src.common import DEVICE, get_dataloader, get_model, train_epoch, validate
+from src.early_stopping import EarlyStopping
 
 
 def main() -> None:
     """Execute the fine-tuning pipeline."""
+    choices = ["baseline", "pipeline1", "pipeline2", "pipeline3"]
     parser = argparse.ArgumentParser()
+    parser.add_argument("--pipeline", choices=choices, type=str, required=True)
     parser.add_argument("--model", type=str, required=True)
-    parser.add_argument("--config", type=str, required=True)
-    parser.add_argument("--output_dir", type=str, required=True)
-
+    # parser.add_argument("--config", type=str, required=True)
+    # parser.add_argument("--output_dir", type=str, required=True)
     args = parser.parse_args()
 
-    with open(args.config) as conf_file:
-        config: dict[str, Any] = yaml.safe_load(conf_file)
+    params_path = Path(args.pipeline) / "params.yaml"
+    with open(params_path) as f:
+        config = yaml.safe_load(f)
 
-    out_dir = Path(args.output_dir)
+    out_dir = Path(args.pipeline / args.model) / "finetuned"
     out_dir.mkdir(parents=True, exist_ok=True)
 
     t_loader = get_dataloader(
@@ -43,6 +45,11 @@ def main() -> None:
     model.load_state_dict(torch.load(weights_path, map_location=DEVICE))
     model.to(DEVICE)
 
+    best_model_path = out_dir / "model.pth"
+    early_stopper = EarlyStopping(
+        alpha=config["train"].get("alpha", 5.0), path=str(best_model_path)
+    )
+
     # Unfreeze layers
     for param in model.parameters():
         param.requires_grad = True
@@ -65,6 +72,15 @@ def main() -> None:
         print(
             f"Epoch {epoch + 1}/{epochs} | T-Loss: {t_loss:.4f} | V-Loss: {v_loss:.4f}"
         )
+
+        early_stopper(v_loss, epoch + 1, model)
+        if early_stopper.stop:
+            print(
+                f"Stopping at epoch {epoch + 1}. "
+                f"Best model was at epoch {early_stopper.best_epoch}"
+            )
+            break
+
     print(f"Model {args.model} fine-tuned successfully!")
 
     # Saving the model

src/preprocessing/augment.py

@@ -107,14 +107,19 @@ def main() -> None:
     """Run data augmentation."""
     parser = argparse.ArgumentParser()
     parser.add_argument("--config", type=str, required=True)
+    parser.add_argument("--force", type=bool, default=False)
     args = parser.parse_args()
 
     with open(args.config) as conf_file:
         config = yaml.safe_load(conf_file)
 
-    process_dataset(
-        Path(config["data"]["inputset_path"]), Path(config["data"]["augmentedset_path"])
-    )
+    input_dir = Path(config["data"]["inputset_path"])
+    output_dir = Path(config["data"]["augmentedset_path"])
+
+    if output_dir.exists() and not args.force:
+        return
+
+    process_dataset(input_dir, output_dir)
 
 
 if __name__ == "__main__":