Repository: autrainer/autrainer

Issue:
Add Tests for BaseMTRegressionDataset
We test all base datasets in `test_datasets.py`, however the corresponding tests for `BaseMTRegressionDataset` are still missing.

Discussion:

Patch:
diff --git a/autrainer/datasets/abstract_dataset.py b/autrainer/datasets/abstract_dataset.py
index 773d350..abf25bb 100644
--- a/autrainer/datasets/abstract_dataset.py
+++ b/autrainer/datasets/abstract_dataset.py
@@ -669,7 +669,7 @@ class BaseMTRegressionDataset(AbstractDataset):
     @cached_property
     def target_transform(self) -> MultiTargetMinMaxScaler:
         return MultiTargetMinMaxScaler(
-            targets=self.target_column,
+            target=self.target_column,
             minimum=self.df_train[self.target_column].min().to_list(),
             maximum=self.df_train[self.target_column].max().to_list(),
         )
diff --git a/autrainer/training/training.py b/autrainer/training/training.py
index b94e50d..79afd86 100644
--- a/autrainer/training/training.py
+++ b/autrainer/training/training.py
@@ -1,7 +1,7 @@
 from copy import deepcopy
 import os
 from pathlib import Path
-from typing import Callable, Dict, List, Tuple
+from typing import Callable, Dict, List, Optional, Tuple
 
 import hydra
 from omegaconf import DictConfig, OmegaConf
@@ -27,7 +27,7 @@ from autrainer.transforms import SmartCompose, TransformManager
 
 from .callback_manager import CallbackManager
 from .continue_training import ContinueTraining
-from .outputs_tracker import init_trackers
+from .outputs_tracker import OutputsTracker, init_trackers
 from .utils import (
     disaggregated_evaluation,
     format_results,
@@ -44,7 +44,7 @@ class ModularTaskTrainer:
         experiment_id: str = None,
         run_name: str = None,
     ) -> None:
-        """Modular Task Trainer.
+        """Trainer managing the training of a model given a configuration.
 
         Args:
             cfg: Run configuration.
@@ -449,6 +449,7 @@ class ModularTaskTrainer:
         ]
 
     def train_epochs(self):
+        """Train the model for a fixed number of epochs."""
         train_loss = []
         self.train_timer.start()
         for epoch in range(self.initial_iteration, self.cfg.iterations + 1):
@@ -531,6 +532,7 @@ class ModularTaskTrainer:
             )
 
     def train_steps(self):
+        """Train the model for a fixed number of steps."""
         pbar = tqdm(
             total=self.cfg.eval_frequency,
             desc="Train",
@@ -633,6 +635,19 @@ class ModularTaskTrainer:
         criterion: torch.nn.Module,
         probabilities_fn: Callable,
     ) -> Tuple[float, torch.Tensor]:
+        """Train the model for a single step.
+
+        Args:
+            model: Model to train.
+            data: Features to train on.
+            target: Target values.
+            criterion: Criterion to optimize.
+            probabilities_fn: Function to convert model output to
+                probabilities.
+
+        Returns:
+            Tuple containing the loss and detached model output.
+        """
         self.optimizer.zero_grad()
         output = model(data)
         loss = criterion(probabilities_fn(output), target)
@@ -643,12 +658,12 @@ class ModularTaskTrainer:
     def evaluate(
         self,
         iteration: int,
-        iteration_folder,
-        loader,
-        df,
-        dev_evaluation=True,
-        save_to="dev",
-        tracker=None,
+        iteration_folder: str,
+        loader: torch.utils.data.DataLoader,
+        df: pd.DataFrame,
+        dev_evaluation: bool = True,
+        save_to: str = "dev",
+        tracker: Optional[OutputsTracker] = None,
     ) -> Dict[str, float]:
         """Evaluate the model on the dev or test set.
 
@@ -762,8 +777,23 @@ class ModularTaskTrainer:
         tracker.reset()
 
     def _evaluate(
-        self, loader, tracker, iteration_folder, cb_type: str = "val"
-    ):
+        self,
+        loader: torch.utils.data.DataLoader,
+        tracker: OutputsTracker,
+        iteration_folder: str,
+        cb_type: str = "val",
+    ) -> Dict[str, float]:
+        """Evaluate the model on the dev or test set.
+
+        Args:
+            loader: Dataloader to evaluate on.
+            tracker: Tracker to save the outputs.
+            iteration_folder: Iteration folder to save the results to.
+            cb_type: Callback type. Defaults to "val".
+
+        Returns:
+            Dictionary containing the results on the dev or test set.
+        """
         with torch.no_grad():
             loss = 0
             for batch_idx, (features, target, sample_idx) in enumerate(
diff --git a/autrainer/training/utils.py b/autrainer/training/utils.py
index 6c7e0d0..2b07734 100644
--- a/autrainer/training/utils.py
+++ b/autrainer/training/utils.py
@@ -37,7 +37,7 @@ def disaggregated_evaluation(
             in stratified manner.
 
     Returns:
-        Dictionary containing results
+        Dictionary containing the results of the disaggregated evaluation.
 
     """
     results = {m.name: {} for m in metrics}
diff --git a/docs/source/modules/training.rst b/docs/source/modules/training.rst
index 0a41691..cd4e505 100644
--- a/docs/source/modules/training.rst
+++ b/docs/source/modules/training.rst
@@ -45,14 +45,14 @@ file and imported in the :ref:`main configuration <main_configuration>` file.
 Trainer
 -------
 
-:class:`autrainer.training.Trainer` manages the training process.
+:class:`~autrainer.training.ModularTaskTrainer` manages the training process.
 It instantiates the model, dataset, criterion, optimizer, scheduler, and callbacks, and trains the model on the dataset.
 It also logs the training process and saves the model, optimizer, and scheduler states at the end of each epoch.
 
 The :attr:`cfg` of the trainer is the composed main configuration file (e.g. :file:`conf/config.yaml`) for each training configuration in the sweep.
 
 .. autoclass:: autrainer.training.ModularTaskTrainer
-   :members:
+   :members: train
 
 
 .. _callbacks:


Test patch:
diff --git a/tests/test_datasets.py b/tests/test_datasets.py
index e0f93ac..6fd80f1 100644
--- a/tests/test_datasets.py
+++ b/tests/test_datasets.py
@@ -13,6 +13,7 @@ from autrainer.datasets import (
     AbstractDataset,
     BaseClassificationDataset,
     BaseMLClassificationDataset,
+    BaseMTRegressionDataset,
     BaseRegressionDataset,
     DCASE2016Task1,
     DCASE2018Task3,
@@ -62,6 +63,7 @@ class TestBaseDatasets(BaseIndividualTempDir):
             "classification",
             "ml-classification",
             "regression",
+            "mt-regression",
         ]
         os.makedirs(os.path.join(path, "default"), exist_ok=True)
         df = pd.DataFrame()
@@ -70,9 +72,12 @@ class TestBaseDatasets(BaseIndividualTempDir):
             df[target_column] = [i % 10 for i in range(num_files)]
         elif target_type == "regression":
             df[target_column] = [i for i in range(num_files)]
-        else:
+        elif target_type == "ml-classification":
             for i in range(10):
                 df[f"target_{i}"] = torch.randint(0, 2, (num_files,)).tolist()
+        else:
+            for i in range(10):
+                df[f"target_{i}"] = torch.rand(num_files).tolist()
 
         output_files = output_files or ["train", "dev", "test"]
         for output_file in output_files:
@@ -204,6 +209,16 @@ class TestBaseDatasets(BaseIndividualTempDir):
         data = BaseRegressionDataset(**kwargs)
         self._test_data(data, 1, (4,))
 
+    def test_mtr(self) -> None:
+        self._mock_dataframes("data/TestDataset", target_type="mt-regression")
+        self._mock_data("data/TestDataset", (101, 64))
+        kwargs = self._mock_dataset_kwargs()
+        kwargs["metrics"] = ["autrainer.metrics.MSE"]
+        kwargs["tracking_metric"] = "autrainer.metrics.MSE"
+        kwargs["target_column"] = [f"target_{i}" for i in range(10)]
+        data = BaseMTRegressionDataset(**kwargs)
+        self._test_data(data, 10, (4, 10))
+
     def _test_data(
         self,
         data: AbstractDataset,