Loss functions module

clinicadl/losses/__init__.py

@@ -0,0 +1,2 @@
+from .config import ClassificationLoss, ImplementedLoss, LossConfig
+from .factory import get_loss_function

clinicadl/losses/config.py

@@ -0,0 +1,85 @@
+from enum import Enum
+from typing import List, Optional, Union
+
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    NonNegativeFloat,
+    NonNegativeInt,
+    PositiveFloat,
+    field_validator,
+)
+
+from clinicadl.utils.enum import BaseEnum
+from clinicadl.utils.factories import DefaultFromLibrary
+
+
+class ClassificationLoss(str, BaseEnum):
+    """Losses that can be used only for classification."""
+
+    CROSS_ENTROPY = "CrossEntropyLoss"
+    MULTI_MARGIN = "MultiMarginLoss"
+
+
+class ImplementedLoss(str, BaseEnum):
+    """Implemented losses in ClinicaDL."""
+
+    CROSS_ENTROPY = "CrossEntropyLoss"
+    MULTI_MARGIN = "MultiMarginLoss"
+    L1 = "L1Loss"
+    MSE = "MSELoss"
+    HUBER = "HuberLoss"
+    SMOOTH_L1 = "SmoothL1Loss"
+
+    @classmethod
+    def _missing_(cls, value):
+        raise ValueError(
+            f"{value} is not implemented. Implemented losses are: "
+            + ", ".join([repr(m.value) for m in cls])
+        )
+
+
+class Reduction(str, Enum):
+    """Supported reduction method in ClinicaDL."""
+
+    NONE = "none"
+    MEAN = "mean"
+    SUM = "sum"
+
+
+class Order(int, Enum):
+    """Supported order of L-norm for MultiMarginLoss."""
+
+    ONE = 1
+    TWO = 2
+
+
+class LossConfig(BaseModel):
+    """Config class to configure the loss function."""
+
+    loss: ImplementedLoss = ImplementedLoss.MSE
+    reduction: Union[Reduction, DefaultFromLibrary] = DefaultFromLibrary.YES
+    delta: Union[PositiveFloat, DefaultFromLibrary] = DefaultFromLibrary.YES
+    beta: Union[PositiveFloat, DefaultFromLibrary] = DefaultFromLibrary.YES
+    p: Union[Order, DefaultFromLibrary] = DefaultFromLibrary.YES
+    margin: Union[float, DefaultFromLibrary] = DefaultFromLibrary.YES
+    weight: Union[
+        Optional[List[NonNegativeFloat]], DefaultFromLibrary
+    ] = DefaultFromLibrary.YES
+    ignore_index: Union[NonNegativeInt, DefaultFromLibrary] = DefaultFromLibrary.YES
+    label_smoothing: Union[
+        NonNegativeFloat, DefaultFromLibrary
+    ] = DefaultFromLibrary.YES
+    # pydantic config
+    model_config = ConfigDict(
+        validate_assignment=True, use_enum_values=True, validate_default=True
+    )
+
+    @field_validator("label_smoothing")
+    @classmethod
+    def validator_label_smoothing(cls, v):
+        if isinstance(v, float):
+            assert (
+                0 <= v <= 1
+            ), f"label_smoothing must be between 0 and 1 but it has been set to {v}."
+        return v

clinicadl/losses/factory.py

@@ -0,0 +1,50 @@
+import inspect
+from copy import deepcopy
+from typing import Any, Dict, Tuple
+
+import torch
+
+from clinicadl.utils.factories import DefaultFromLibrary, get_args_and_defaults
+
+from .config import LossConfig
+
+
+def get_loss_function(config: LossConfig) -> Tuple[torch.nn.Module, Dict[str, Any]]:
+    """
+    Factory function to get a loss function from PyTorch.
+
+    Parameters
+    ----------
+    loss : LossConfig
+        The config class with the parameters of the loss function.
+
+    Returns
+    -------
+    nn.Module
+        The loss function.
+    Dict[str, Any]
+        The config dict with only the parameters relevant to the selected
+        loss function.
+    """
+    loss_class = getattr(torch.nn, config.loss)
+    expected_args, config_dict = get_args_and_defaults(loss_class.__init__)
+    for arg, value in config.model_dump().items():
+        if arg in expected_args and value != DefaultFromLibrary.YES:
+            config_dict[arg] = value
+
+    config_dict_ = deepcopy(config_dict)
+    if "weight" in config_dict and config_dict["weight"] is not None:
+        config_dict_["weight"] = torch.Tensor(config_dict_["weight"])
+    loss = loss_class(**config_dict_)
+
+    config_dict["loss"] = config.loss
+
+    return loss, config_dict
+
+
+# TODO : what about them?
+# "KLDivLoss",
+# "BCEWithLogitsLoss",
+# "VAEGaussianLoss",
+# "VAEBernoulliLoss",
+# "VAEContinuousBernoulliLoss",

clinicadl/utils/factories.py

@@ -0,0 +1,34 @@
+import inspect
+from enum import Enum
+from typing import Any, Callable, Dict, List, Tuple
+
+
+class DefaultFromLibrary(str, Enum):
+    YES = "DefaultFromLibrary"
+
+
+def get_args_and_defaults(func: Callable) -> Tuple[List[str], Dict[str, Any]]:
+    """
+    Gets the arguments of a function, as well as the default
+    values possibly attached to them.
+
+    Parameters
+    ----------
+    func : Callable
+        The function.
+
+    Returns
+    -------
+    List[str]
+        The names of the arguments.
+    Dict[str, Any]
+        The default values in a dict.
+    """
+    signature = inspect.signature(func)
+    args = list(signature.parameters.keys())
+    defaults = {
+        k: v.default
+        for k, v in signature.parameters.items()
+        if v.default is not inspect.Parameter.empty
+    }
+    return args, defaults

tests/unittests/losses/test_config.py

@@ -0,0 +1,17 @@
+import pytest
+from pydantic import ValidationError
+
+
+def test_LossConfig():
+    from clinicadl.losses import LossConfig
+
+    LossConfig(reduction="none", p=2, weight=[0.1, 0.1, 0.8])
+    LossConfig(
+        loss="SmoothL1Loss", margin=10.0, delta=2.0, beta=3.0, label_smoothing=0.5
+    )
+    with pytest.raises(ValueError):
+        LossConfig(loss="abc")
+    with pytest.raises(ValueError):
+        LossConfig(weight=[0.1, -0.1, 0.8])
+    with pytest.raises(ValidationError):
+        LossConfig(label_smoothing=1.1)

tests/unittests/losses/test_factory.py

@@ -0,0 +1,27 @@
+from torch import Tensor
+from torch.nn import MultiMarginLoss
+
+from clinicadl.losses import ImplementedLoss, LossConfig, get_loss_function
+
+
+def test_get_loss_function():
+    for loss in [e.value for e in ImplementedLoss]:
+        config = LossConfig(loss=loss)
+        get_loss_function(config)
+
+    config = LossConfig(loss="MultiMarginLoss", reduction="sum", weight=[1, 2, 3], p=2)
+    loss, config_dict = get_loss_function(config)
+    assert isinstance(loss, MultiMarginLoss)
+    assert loss.reduction == "sum"
+    assert loss.p == 2
+    assert loss.margin == 1.0
+    assert (loss.weight == Tensor([1, 2, 3])).all()
+    assert config_dict == {
+        "loss": "MultiMarginLoss",
+        "reduction": "sum",
+        "p": 2,
+        "margin": 1.0,
+        "weight": [1, 2, 3],
+        "size_average": None,
+        "reduce": None,
+    }

tests/unittests/utils/test_factories_utils.py

@@ -0,0 +1,10 @@
+from clinicadl.utils.factories import get_args_and_defaults
+
+
+def test_get_default_args():
+    def f(a, b="b", c=0, d=None):
+        return None
+
+    args, defaults = get_args_and_defaults(f)
+    assert args == ["a", "b", "c", "d"]
+    assert defaults == {"b": "b", "c": 0, "d": None}