Implemented alternative experiment tracking functionality (#102)

.gitignore

@@ -111,8 +111,9 @@ data/
 # Models
 models/
 
-# Weights and Biases experiment tracking
+# Experiment tracking
 wandb/
+mlruns/
 
 # Data files
 *.xlsx

README.md

@@ -13,7 +13,10 @@ ______________________________________________________________________
 
 Developers:
 
+- Anders Jess Pedersen ([email protected])
 - Dan Saattrup Nielsen ([email protected])
+- Simon Leminen Madsen ([email protected])
+
 
 
 ## Installation

config/asr_finetuning.yaml

@@ -1,16 +1,19 @@
 defaults:
-  - model: wav2vec2-small
+  - model: whisper-xxsmall
   - datasets:
     - coral
   - decoder_datasets:
     - wikipedia
     - common_voice
     - reddit
+  - experiment_tracking: wandb
   - override hydra/job_logging: custom
   - _self_
 
 seed: 4242
 
+experiment_tracking: null
+
 evaluation_dataset:
   id: alexandrainst/coral
   subset: read_aloud
@@ -48,10 +51,6 @@ fp16_allowed: true
 bf16_allowed: true
 
 # Training parameters
-wandb: false
-wandb_project: CoRal
-wandb_group: default
-wandb_name: ${model_id}
 resume_from_checkpoint: false
 ignore_data_skip: false
 save_total_limit: 0  # Will automatically be set to >=1 if `early_stopping` is enabled

config/experiment_tracking/mlflow.yaml

@@ -0,0 +1,3 @@
+type: mlflow
+name_experiment: CoRal
+name_run: ${model_id}

config/experiment_tracking/wandb.yaml

@@ -0,0 +1,4 @@
+type: wandb
+name_experiment: CoRal
+name_run: ${model_id}
+name_group: default

pyproject.toml

@@ -36,6 +36,7 @@ gradio = {version = "^5.5.0", optional=true}
 samplerate = {version="^0.2.1", optional=true}
 punctfix = {version="^0.11.1", optional=true}
 matplotlib = {version = "^3.9.2", optional = true}
+mlflow = "^2.17.2"
 
 [tool.poetry.group.dev.dependencies]
 pytest = ">=8.1.1"

src/coral/experiment_tracking/__init__.py

@@ -0,0 +1,4 @@
+"""The CoRal project.
+
+Experiment tracking.
+"""

src/coral/experiment_tracking/extracking_factory.py

@@ -0,0 +1,28 @@
+"""Factory for experiment tracking setup."""
+
+from omegaconf import DictConfig
+
+from .extracking_setup import ExTrackingSetup
+from .mlflow_setup import MLFlowSetup
+from .wandb_setup import WandbSetup
+
+
+def load_extracking_setup(config: DictConfig) -> ExTrackingSetup:
+    """Return the experiment tracking setup.
+
+    Args:
+        config:
+            The configuration object.
+
+    Returns:
+        The experiment tracking setup.
+    """
+    match config.experiment_tracking.type:
+        case "wandb":
+            return WandbSetup(config=config)
+        case "mlflow":
+            return MLFlowSetup(config=config)
+        case _:
+            raise ValueError(
+                f"Unknown experiment tracking type: {config.experiment_tracking.type}"
+            )

src/coral/experiment_tracking/extracking_setup.py

@@ -0,0 +1,34 @@
+"""This module contains the base class for an experiment tracking setup."""
+
+from abc import ABC, abstractmethod
+
+from omegaconf import DictConfig
+
+
+class ExTrackingSetup(ABC):
+    """Base class for an experiment tracking setup."""
+
+    @abstractmethod
+    def __init__(self, config: DictConfig) -> None:
+        """Initialise the experiment tracking setup.
+
+        Args:
+            config:
+                The configuration object.
+        """
+
+    @abstractmethod
+    def run_initialization(self) -> None:
+        """Run the initialization of the experiment tracking setup.
+
+        Returns:
+            True if the initialization was successful, False otherwise.
+        """
+
+    @abstractmethod
+    def run_finalization(self) -> None:
+        """Run the finalization of the experiment tracking setup.
+
+        Returns:
+            True if the finalization was successful, False otherwise.
+        """

src/coral/experiment_tracking/mlflow_setup.py

@@ -0,0 +1,33 @@
+"""MLFlow experiment tracking setup class."""
+
+import os
+
+import mlflow
+from omegaconf import DictConfig
+
+from .extracking_setup import ExTrackingSetup
+
+
+class MLFlowSetup(ExTrackingSetup):
+    """MLFlow setup class."""
+
+    def __init__(self, config: DictConfig) -> None:
+        """Initialise the MLFlow setup.
+
+        Args:
+            config:
+                The configuration object.
+        """
+        self.config = config
+        self.is_main_process = os.getenv("RANK", "0") == "0"
+
+    def run_initialization(self) -> None:
+        """Run the initialization of the experiment tracking setup."""
+        mlflow.set_experiment(self.config.experiment_tracking.name_experiment)
+        mlflow.start_run(run_name=self.config.experiment_tracking.name_run)
+        return
+
+    def run_finalization(self) -> None:
+        """Run the finalization of the experiment tracking setup."""
+        mlflow.end_run()
+        return

src/coral/experiment_tracking/wandb_setup.py

@@ -0,0 +1,37 @@
+"""wandb experiment tracking setup class."""
+
+import os
+
+import wandb
+from omegaconf import DictConfig
+
+from .extracking_setup import ExTrackingSetup
+
+
+class WandbSetup(ExTrackingSetup):
+    """Wandb setup class."""
+
+    def __init__(self, config: DictConfig) -> None:
+        """Initialise the Wandb setup.
+
+        Args:
+            config:
+                The configuration object.
+        """
+        self.config = config
+        self.is_main_process = os.getenv("RANK", "0") == "0"
+
+    def run_initialization(self) -> None:
+        """Run the initialization of the experiment tracking setup."""
+        wandb.init(
+            project=self.config.experiment_tracking.name_experiment,
+            name=self.config.experiment_tracking.name_run,
+            group=self.config.experiment_tracking.name_group,
+            config=dict(self.config),
+        )
+        return
+
+    def run_finalization(self) -> None:
+        """Run the finalization of the experiment tracking setup."""
+        wandb.finish()
+        return

src/coral/finetune.py

@@ -5,11 +5,10 @@
 
 from omegaconf import DictConfig
 from transformers import EarlyStoppingCallback, TrainerCallback
-from wandb import finish as wandb_finish
-from wandb.sdk.wandb_init import init as wandb_init
 
 from .data import load_data_for_finetuning
 from .data_models import ModelSetup
+from .experiment_tracking.extracking_factory import load_extracking_setup
 from .model_setup import load_model_setup
 from .ngram import train_and_store_ngram_model
 from .utils import block_terminal_output, disable_tqdm, push_model_to_hub
@@ -33,13 +32,9 @@ def finetune(config: DictConfig) -> None:
     model = model_setup.load_model()
     dataset = load_data_for_finetuning(config=config, processor=processor)
 
-    if config.wandb and is_main_process:
-        wandb_init(
-            project=config.wandb_project,
-            group=config.wandb_group,
-            name=config.wandb_name,
-            config=dict(config),
-        )
+    if bool(config.experiment_tracking) and is_main_process:
+        extracking_setup = load_extracking_setup(config=config)
+        extracking_setup.run_initialization()
 
     if "val" not in dataset and is_main_process:
         logger.info("No validation set found. Disabling early stopping.")
@@ -58,8 +53,9 @@ def finetune(config: DictConfig) -> None:
     block_terminal_output()
     with disable_tqdm():
         trainer.train(resume_from_checkpoint=config.resume_from_checkpoint)
-    if config.wandb and is_main_process:
-        wandb_finish()
+
+    if bool(config.experiment_tracking) and is_main_process:
+        extracking_setup.run_finalization()
 
     model.save_pretrained(save_directory=config.model_dir)
 

src/coral/wav2vec2.py

@@ -206,7 +206,9 @@ def load_training_arguments(self) -> TrainingArguments:
             optim=OptimizerNames.ADAMW_TORCH,
             adam_beta1=self.config.adam_first_momentum,
             adam_beta2=self.config.adam_second_momentum,
-            report_to=["wandb"] if self.config.wandb else [],
+            report_to=[self.config.experiment_tracking.type]
+            if self.config.experiment_tracking
+            else [],
             ignore_data_skip=self.config.ignore_data_skip,
             save_safetensors=True,
             use_cpu=hasattr(sys, "_called_from_test"),

src/coral/whisper.py

@@ -190,7 +190,9 @@ def load_training_arguments(self) -> TrainingArguments:
             optim=OptimizerNames.ADAMW_TORCH,
             adam_beta1=self.config.adam_first_momentum,
             adam_beta2=self.config.adam_second_momentum,
-            report_to=["wandb"] if self.config.wandb else [],
+            report_to=[self.config.experiment_tracking.type]
+            if self.config.experiment_tracking
+            else [],
             ignore_data_skip=self.config.ignore_data_skip,
             save_safetensors=True,
             predict_with_generate=True,