Repository: autrainer/autrainer

Issue:
Per-Example Loss Values are Wrong for BalancedCrossEntropyLoss
In our current implementation of the `BalancedCrossEntropyLoss` we calculate the class weights based on the inverse frequency of each class in the training subset, which is correct. However, we have an oversight that we normalize the weights to sum up to `1` instead of the number of classes `C`:
https://github.com/autrainer/autrainer/blob/c25132ee008508fb9cbd941724c6e9bfe86a0a4b/autrainer/criterions/criterion_wrappers.py#L28-L45

This is no issue when setting `reduction="mean"` which is the default for the [CrossEntropyLoss](https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html#torch.nn.CrossEntropyLoss) implementation, as torch automatically rescales the weights in this case (can also be seen in the formulas and [here](https://github.com/pytorch/pytorch/blob/2981534f54d49fa3a9755c9b0855e7929c2527f0/aten/src/ATen/native/LossNLL.cpp#L265-L267)).

**BUT** when obtaining the per-example loss again in our `OutputsTracker`, we end up with loss values that are essentially downscaled by a factor of `C`.
This is not "severe" in a sense that the loss computation is wrong during training, but all our reported per-example loss values are definitely not correct.

Discussion:

Patch:
diff --git a/autrainer/core/wrappers/instantiate.py b/autrainer/core/wrappers/instantiate.py
index f857da5..800e953 100644
--- a/autrainer/core/wrappers/instantiate.py
+++ b/autrainer/core/wrappers/instantiate.py
@@ -36,7 +36,7 @@ def instantiate(
     convert: Optional[HydraConvertEnum] = None,
     recursive: bool = False,
     **kwargs,
-) -> Optional[T]:
+) -> T:
     """Instantiate an object from a configuration Dict or DictConfig.
 
     The config must contain a `_target_` field that specifies a relative import
@@ -94,7 +94,7 @@ def instantiate_shorthand(
     convert: Optional[HydraConvertEnum] = None,
     recursive: bool = False,
     **kwargs,
-) -> Optional[T]:
+) -> T:
     """Instantiate an object from a shorthand configuration.
 
     A shorthand config is either a string or a dictionary with a single key.
diff --git a/autrainer/criterions/criterion_wrappers.py b/autrainer/criterions/criterion_wrappers.py
index c0e37a0..cb98c15 100644
--- a/autrainer/criterions/criterion_wrappers.py
+++ b/autrainer/criterions/criterion_wrappers.py
@@ -41,8 +41,7 @@ class BalancedCrossEntropyLoss(CrossEntropyLoss):
             .values
         )
         weight = torch.tensor(1 / frequency, dtype=torch.float32)
-        weight /= weight.sum()
-        self.weight = weight
+        self.weight = weight * len(weight) / weight.sum()
 
 
 class MSELoss(torch.nn.MSELoss):
diff --git a/autrainer/optimizers/sam.py b/autrainer/optimizers/sam.py
index 4a54bd1..f83b25f 100644
--- a/autrainer/optimizers/sam.py
+++ b/autrainer/optimizers/sam.py
@@ -139,7 +139,7 @@ class SAM(torch.optim.Optimizer):
         target: torch.Tensor,
         criterion: torch.nn.Module,
         probabilities_fn: Callable,
-    ) -> Tuple[float, torch.Tensor]:
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
         """Sharpness Aware Minimization requires two forward-backward passes
         over the batch to simultaneously minimize the loss value and the loss
         sharpness.
@@ -152,16 +152,15 @@ class SAM(torch.optim.Optimizer):
             probabilities_fn: Function to get probabilities from model outputs.
 
         Returns:
-            Reduced loss over the batch and detached model outputs.
+            Tuple containing the non-reduced loss and model outputs.
         """
         output = model(data)
         loss = criterion(probabilities_fn(output), target)
-        loss.backward()
+        loss.mean().backward()
         self.first_step(zero_grad=True)
 
         output = model(data)
         loss = criterion(probabilities_fn(output), target)
-        loss.backward()
+        loss.mean().backward()
         self.second_step(zero_grad=True)
-        _loss = loss.item()
-        return _loss, output.detach()
+        return loss, output
diff --git a/autrainer/training/outputs_tracker.py b/autrainer/training/outputs_tracker.py
index dfa02b8..6e34757 100644
--- a/autrainer/training/outputs_tracker.py
+++ b/autrainer/training/outputs_tracker.py
@@ -1,8 +1,6 @@
-import copy
-from typing import List
+from typing import Callable, List, TypeVar
 
 import numpy as np
-from omegaconf import DictConfig
 import pandas as pd
 import torch
 
@@ -10,28 +8,41 @@ from autrainer.core.utils import Bookkeeping
 from autrainer.datasets import AbstractDataset
 
 
+T = TypeVar("T")
+
+
 class OutputsTracker:
     def __init__(
         self,
         export: bool,
         prefix: str,
         data: AbstractDataset,
-        criterion: DictConfig,
-        bookkeeping: Bookkeeping = None,
+        bookkeeping: Bookkeeping,
     ) -> None:
+        """Tracker for model outputs, targets, losses, and predictions.
+
+        Args:
+            export: Whether to export the results.
+            prefix: Prefix for the exported files.
+            data: Instance of the dataset.
+            bookkeeping: Instance of the bookkeeping class.
+        """
         self._export = export
         self._prefix = prefix
         self._data = data
         self._bookkeeping = bookkeeping
-        self._tracker_criterion = copy.deepcopy(criterion).cpu()
-        self._tracker_criterion.reduction = "none"
         self.reset()
 
     def reset(self) -> None:
+        """Reset the tracker. Clears all the stored data accumulated over a
+        single iteration.
+        """
         self._outputs = torch.zeros(
             (0, self._data.output_dim), dtype=torch.float32
         )
-        self._targets = torch.zeros(0, dtype=torch.long)
+        self._targets = torch.zeros(
+            0, dtype=torch.long
+        )  # automatic type promotion in case of float targets
         self._indices = torch.zeros(0, dtype=torch.long)
         self._losses = torch.zeros(0, dtype=torch.float32)
         self._predictions = None
@@ -41,19 +52,31 @@ class OutputsTracker:
         self,
         output: torch.Tensor,
         target: torch.Tensor,
+        loss: torch.Tensor,
         sample_idx: torch.Tensor,
     ) -> None:
-        output = output.cpu()
-        target = target.cpu()
-        self._outputs = torch.cat([self._outputs, output], dim=0)
-        self._targets = torch.cat([self._targets, target], dim=0)
-        self._indices = torch.cat([self._indices, sample_idx], dim=0)
-
-        with torch.no_grad():
-            loss = self._tracker_criterion(output, target)
-            self._losses = torch.cat([self._losses, loss], dim=0)
-
-    def save(self, iteration_folder: str, reset=True) -> None:
+        """Update the tracker with the current model outputs, targets, losses,
+        and sample indices.
+
+        Args:
+            output: Detached model outputs.
+            target: Targets.
+            loss: Per-sample losses.
+            sample_idx: Sample indices.
+        """
+        self._outputs = torch.cat([self._outputs, output.cpu()], dim=0)
+        self._targets = torch.cat([self._targets, target.cpu()], dim=0)
+        self._losses = torch.cat([self._losses, loss.cpu()], dim=0)
+        self._indices = torch.cat([self._indices, sample_idx.cpu()], dim=0)
+
+    def save(self, iteration_folder: str, reset: bool = True) -> None:
+        """Save the tracked data to disk.
+
+        Args:
+            iteration_folder: Current iteration folder.
+            reset: Whether to reset the tracker after saving the results.
+                Defaults to True.
+        """
         results = {
             "outputs": self._outputs.numpy(),
             "targets": self._targets.numpy(),
@@ -95,8 +118,8 @@ class OutputsTracker:
         if reset:
             self.reset()
 
-    def check_saved(func):
-        def wrapper(self, *args, **kwargs):
+    def check_saved(func: Callable[..., T]) -> Callable[..., T]:
+        def wrapper(self: "OutputsTracker", *args, **kwargs) -> T:
             if self._results_df is None:
                 raise ValueError("Results not saved yet.")
             return func(self, *args, **kwargs)
@@ -106,31 +129,61 @@ class OutputsTracker:
     @property
     @check_saved
     def outputs(self) -> np.ndarray:
+        """Get the model outputs.
+
+        Returns:
+            Model outputs.
+        """
         return self._outputs.numpy()
 
     @property
     @check_saved
     def targets(self) -> np.ndarray:
+        """Get the targets.
+
+        Returns:
+            Targets.
+        """
         return self._targets.numpy()
 
     @property
     @check_saved
     def indices(self) -> np.ndarray:
+        """Get the sample indices.
+
+        Returns:
+            Sample indices.
+        """
         return self._indices.numpy()
 
     @property
     @check_saved
     def losses(self) -> np.ndarray:
+        """Get the per-sample losses.
+
+        Returns:
+            Per-sample losses.
+        """
         return self._losses.numpy()
 
     @property
     @check_saved
     def predictions(self) -> np.ndarray:
+        """Get the predictions.
+
+        Returns:
+            Predictions.
+        """
         return np.array(self._predictions)
 
     @property
     @check_saved
     def results_df(self) -> pd.DataFrame:
+        """Get the results as a DataFrame.
+
+        Returns:
+            Results as a DataFrame.
+        """
         return self._results_df
 
 
@@ -138,9 +191,19 @@ def init_trackers(
     exports: List[bool],
     prefixes: List[str],
     data: AbstractDataset,
-    criterion: DictConfig,
-    bookkeeping: Bookkeeping = None,
+    bookkeeping: Bookkeeping,
 ) -> List[OutputsTracker]:
+    """Utility function to initialize multiple trackers at once.
+
+    Args:
+        exports: Whether to export the results.
+        prefixes: Prefixes for the exported files.
+        data: Instance of the dataset.
+        bookkeeping: Instance of the bookkeeping class.
+
+    Returns:
+        List of initialized trackers.
+    """
     trackers = []
     for export, prefix in zip(exports, prefixes):
         trackers.append(
@@ -148,7 +211,6 @@ def init_trackers(
                 export=export,
                 prefix=prefix,
                 data=data,
-                criterion=criterion,
                 bookkeeping=bookkeeping,
             )
         )
diff --git a/autrainer/training/training.py b/autrainer/training/training.py
index c69b0ad..4a8d237 100644
--- a/autrainer/training/training.py
+++ b/autrainer/training/training.py
@@ -121,6 +121,7 @@ class ModularTaskTrainer:
         self.criterion = autrainer.instantiate_shorthand(
             config=self.cfg.criterion,
             instance_of=torch.nn.modules.loss._Loss,
+            reduction="none",
         )
         if hasattr(self.criterion, "setup"):
             self.criterion.setup(self.data)
@@ -336,17 +337,15 @@ class ModularTaskTrainer:
             exports=[self.cfg.save_dev_outputs, self.cfg.save_test_outputs],
             prefixes=["dev", "test"],
             data=self.data,
-            criterion=self.criterion,
             bookkeeping=self.bookkeeping,
         )
         if self.cfg.save_train_outputs:
-            self.train_tracker = init_trackers(
-                exports=[True],
-                prefixes=["train"],
+            self.train_tracker = OutputsTracker(
+                export=True,
+                prefix="train",
                 data=self.data,
-                criterion=self.criterion,
                 bookkeeping=self.bookkeeping,
-            )[0]
+            )
         else:
             self.train_tracker = None
 
@@ -470,25 +469,33 @@ class ModularTaskTrainer:
                     iteration=epoch,
                     batch_idx=batch_idx,
                 )
-                l, o = self.train_step_fn(
+                loss, output = self.train_step_fn(
                     self.model,
                     data,
                     target,
                     self.criterion,
                     self.data.target_transform.probabilities_training,
                 )
+                loss = loss.detach()
+                output = output.detach()
+                reduced_loss = loss.mean().item()
                 if self.scheduler and self.scheduler_frequency == "batch":
                     self.scheduler.step()
                 if self.train_tracker:
-                    self.train_tracker.update(o, target, sample_idx)
+                    self.train_tracker.update(
+                        output,
+                        target,
+                        loss,
+                        sample_idx,
+                    )
                 self.callback_manager.callback(
                     position="cb_on_step_end",
                     trainer=self,
                     iteration=epoch,
                     batch_idx=batch_idx,
-                    loss=l,
+                    loss=reduced_loss,
                 )
-                train_loss.append(l)
+                train_loss.append(reduced_loss)
             if epoch % self.cfg.eval_frequency == 0:
                 if self.scheduler and self.scheduler_frequency == "evaluation":
                     self.scheduler.step()
@@ -563,25 +570,28 @@ class ModularTaskTrainer:
                 iteration=step,
                 batch_idx=(step - 1) % self.cfg.eval_frequency,
             )
-            l, o = self.train_step_fn(
+            loss, output = self.train_step_fn(
                 self.model,
                 data,
                 target,
                 self.criterion,
                 self.data.target_transform.probabilities_training,
             )
+            loss = loss.detach()
+            output = output.detach()
+            reduced_loss = loss.mean().item()
             if self.scheduler and self.scheduler_frequency == "batch":
                 self.scheduler.step()
             if self.train_tracker:
-                self.train_tracker.update(o, target, sample_idx)
+                self.train_tracker.update(output, target, loss, sample_idx)
             self.callback_manager.callback(
                 position="cb_on_step_end",
                 trainer=self,
                 iteration=step,
                 batch_idx=(step - 1) % self.cfg.eval_frequency,
-                loss=l,
+                loss=reduced_loss,
             )
-            train_loss.append(l)
+            train_loss.append(reduced_loss)
             if step % self.cfg.eval_frequency == 0:
                 if self.scheduler and self.scheduler_frequency == "evaluation":
                     self.scheduler.step()
@@ -630,7 +640,7 @@ class ModularTaskTrainer:
         target: torch.Tensor,
         criterion: torch.nn.Module,
         probabilities_fn: Callable,
-    ) -> Tuple[float, torch.Tensor]:
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
         """Train the model for a single step.
 
         Args:
@@ -642,14 +652,14 @@ class ModularTaskTrainer:
                 probabilities.
 
         Returns:
-            Tuple containing the loss and detached model output.
+            Tuple containing the non-reduced loss and model outputs.
         """
         self.optimizer.zero_grad()
         output = model(data)
         loss = criterion(probabilities_fn(output), target)
-        loss.backward()
+        loss.mean().backward()
         self.optimizer.step()
-        return loss.item(), output.detach()
+        return loss, output
 
     def evaluate(
         self,
@@ -791,7 +801,7 @@ class ModularTaskTrainer:
             Dictionary containing the results on the dev or test set.
         """
         with torch.no_grad():
-            loss = 0
+            losses = 0
             for batch_idx, (features, target, sample_idx) in enumerate(
                 tqdm(
                     loader,
@@ -805,28 +815,23 @@ class ModularTaskTrainer:
                     batch_idx=batch_idx,
                 )
                 output = self.model(features.to(self.DEVICE))
-                probabilities_fn = (
-                    self.data.target_transform.probabilities_training
-                )
-                loss += (
-                    self.criterion(
-                        probabilities_fn(output),
-                        target.to(self.DEVICE),
-                    )
-                    .cpu()
-                    .item()
+                loss = self.criterion(
+                    self.data.target_transform.probabilities_training(output),
+                    target.to(self.DEVICE),
                 )
-                tracker.update(output, target, sample_idx)
+                reduced_loss = loss.mean().item()
+                losses += reduced_loss
+                tracker.update(output, target, loss, sample_idx)
                 self.callback_manager.callback(
                     position=f"cb_on_{cb_type}_step_end",
                     trainer=self,
                     batch_idx=batch_idx,
-                    loss=loss,
+                    loss=reduced_loss,
                 )
-            loss /= len(loader) + 1
+            losses /= len(loader) + 1
         tracker.save(iteration_folder, reset=False)
         results = {
-            "loss": loss,
+            "loss": losses,
         }
         for metric in self.data.metrics:
             results[metric.name] = metric(tracker.targets, tracker.predictions)
diff --git a/docs/source/examples/optimizer_custom_step.py b/docs/source/examples/optimizer_custom_step.py
index 4ebaec0..53277c9 100644
--- a/docs/source/examples/optimizer_custom_step.py
+++ b/docs/source/examples/optimizer_custom_step.py
@@ -11,7 +11,7 @@ class SomeOptimizer(torch.optim.Optimizer):
         target: torch.Tensor,
         criterion: torch.nn.Module,
         probabilities_fn: Callable,
-    ) -> Tuple[float, torch.Tensor]:
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
         """Custom step function for the optimizer.
 
         Args:
@@ -22,5 +22,5 @@ class SomeOptimizer(torch.optim.Optimizer):
             probabilities_fn: Function to get probabilities from model outputs.
 
         Returns:
-            Reduced loss over the batch and detached model outputs.
+            Tuple containing the non-reduced loss and model outputs.
         """
diff --git a/docs/source/examples/tutorials/random_scaled_sgd.py b/docs/source/examples/tutorials/random_scaled_sgd.py
index 6c4741c..0e91970 100644
--- a/docs/source/examples/tutorials/random_scaled_sgd.py
+++ b/docs/source/examples/tutorials/random_scaled_sgd.py
@@ -35,13 +35,13 @@ class RandomScaledSGD(torch.optim.Optimizer):
         target: torch.Tensor,  # batched target data
         criterion: torch.nn.Module,  # loss function
         probabilities_fn: Callable,  # function to get probabilities from model outputs
-    ) -> Tuple[float, torch.Tensor]:
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
         self.zero_grad()
         output = model(data)
         loss = criterion(probabilities_fn(output), target)
-        loss.backward()
+        loss.mean().backward()
         if torch.rand(1, generator=self.g).item() < self.p:
             self.param_groups[0]["lr"] *= self.scaling_factor
         self.step()
         self.param_groups[0]["lr"] = self.base_lr
-        return loss.item(), output.detach()
+        return loss, output
diff --git a/docs/source/modules/criterions.rst b/docs/source/modules/criterions.rst
index c0ecc7d..a1ffa70 100644
--- a/docs/source/modules/criterions.rst
+++ b/docs/source/modules/criterions.rst
@@ -10,6 +10,14 @@ Criterions are specified in the :ref:`Dataset<datasets>` configuration with :ref
    To create custom criterions, refer to the :ref:`custom criterions tutorial <tut_criterions>`.
 
 
+.. note::
+
+   The :attr:`reduction` attribute of each criterion is automatically set to :attr:`"none"` during training, validation, and testing.
+   This allows the per-example loss to be reported directly, without the need for re-calculating the loss for logging purposes.
+
+   This is handled automatically during the instantiation of the criterion.
+
+
 Criterion Wrappers
 ------------------
 
diff --git a/docs/source/modules/optimizers.rst b/docs/source/modules/optimizers.rst
index 41a9038..bb7d2ab 100644
--- a/docs/source/modules/optimizers.rst
+++ b/docs/source/modules/optimizers.rst
@@ -31,9 +31,8 @@ For example, :class:`torch.optim.Adam` can be used as follows:
 
    .. configurations::
       :subdir: optimizer
-      :configs: Adam SGD
-      :headline:
-
+      :configs: SGD Adam AdamW Adadelta Adagrad Adamax NAdam RAdam SparseAdam RMSprop Rprop ASGD LBFGS
+      :exact:
 
 .. _custom_optimizers:
 
diff --git a/docs/source/usage/tutorials.rst b/docs/source/usage/tutorials.rst
index 96f68f7..ee7e304 100644
--- a/docs/source/usage/tutorials.rst
+++ b/docs/source/usage/tutorials.rst
@@ -194,6 +194,11 @@ To create a custom :ref:`criterion <criterions>`, inherit from :class:`torch.nn.
 If the criterion relies on the dataset, an optional :ref:`criterion setup method <criterion_setup>`
 can be defined which is called after the dataset is initialized.
 
+.. note::
+
+   The :attr:`reduction` attribute of each criterion is automatically set to :attr:`"none"` during instantiation and
+   the :meth:`forward` method should return the per-example loss.
+
 For example, the following criterion implements `CrossEntropyLoss <https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html>`_
 with an additional scaling factor:
 


Test patch:
diff --git a/tests/test_criterions.py b/tests/test_criterions.py
index 1817b75..a33537a 100644
--- a/tests/test_criterions.py
+++ b/tests/test_criterions.py
@@ -27,7 +27,7 @@ class TestBalancedCrossEntropyLoss:
             criterion.weight is not None
         ), "Should have calculated the weights"
         test_weights = torch.tensor([1 / 1, 1 / 2, 1 / 3], dtype=torch.float32)
-        test_weights /= test_weights.sum()
+        test_weights = test_weights * len(test_weights) / test_weights.sum()
         assert torch.allclose(
             criterion.weight, test_weights
         ), "Should have calculated frequency-based weights"