careamics 0.1.0rc2__py3-none-any.whl → 0.1.0rc3__py3-none-any.whl

careamics/config/__init__.py

@@ -1,11 +1,35 @@
 """Configuration module."""
 
 
-__all__ = ["Configuration", "load_configuration", "save_configuration"]
+__all__ = [
+    "AlgorithmModel",
+    "DataModel",
+    "Configuration",
+    "CheckpointModel",
+    "InferenceModel",
+    "load_configuration",
+    "save_configuration",
+    "TrainingModel",
+    "create_n2v_configuration",
+    "register_model",
+    "CustomModel",
+    "create_inference_configuration",
+    "clear_custom_models",
+    "ConfigurationInformation",
+]
 
-from .config import (
+from .algorithm_model import AlgorithmModel
+from .architectures import CustomModel, clear_custom_models, register_model
+from .callback_model import CheckpointModel
+from .configuration_factory import (
+    create_inference_configuration,
+    create_n2v_configuration,
+)
+from .configuration_model import (
     Configuration,
     load_configuration,
     save_configuration,
 )
-from .torch_optim import get_parameters as get_parameters
+from .data_model import DataModel
+from .inference_model import InferenceModel
+from .training_model import TrainingModel

careamics/config/algorithm_model.py

@@ -0,0 +1,167 @@
+from __future__ import annotations
+
+from pprint import pformat
+from typing import Literal, Union
+
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+from typing_extensions import Self
+
+from .architectures import CustomModel, UNetModel, VAEModel
+from .optimizer_models import LrSchedulerModel, OptimizerModel
+
+
+class AlgorithmModel(BaseModel):
+    """Algorithm configuration.
+
+    This Pydantic model validates the parameters governing the components of the
+    training algorithm: which algorithm, loss function, model architecture, optimizer,
+    and learning rate scheduler to use.
+
+    Currently, we only support N2V and custom algorithms. The `n2v` algorithm is only
+    compatible with `n2v` loss and `UNet` architecture. The `custom` algorithm allows
+    you to register your own architecture and select it using its name as
+    `name` in the custom pydantic model.
+
+    Attributes
+    ----------
+    algorithm : Literal["n2v", "custom"]
+        Algorithm to use.
+    loss : Literal["n2v", "mae", "mse"]
+        Loss function to use.
+    model : Union[UNetModel, VAEModel, CustomModel]
+        Model architecture to use.
+    optimizer : OptimizerModel, optional
+        Optimizer to use.
+    lr_scheduler : LrSchedulerModel, optional
+        Learning rate scheduler to use.
+
+    Raises
+    ------
+    ValueError
+        Algorithm parameter type validation errors.
+    ValueError
+        If the algorithm, loss and model are not compatible.
+
+    Examples
+    --------
+    Minimum example:
+    >>> from careamics.config import AlgorithmModel
+    >>> config_dict = {
+    ...     "algorithm": "n2v",
+    ...     "loss": "n2v",
+    ...     "model": {
+    ...         "architecture": "UNet",
+    ...     }
+    ... }
+    >>> config = AlgorithmModel(**config_dict)
+
+    Using a custom model:
+    >>> from torch import nn, ones
+    >>> from careamics.config import AlgorithmModel, register_model
+    ...
+    >>> @register_model(name="linear_model")
+    ... class LinearModel(nn.Module):
+    ...    def __init__(self, in_features, out_features, *args, **kwargs):
+    ...        super().__init__()
+    ...        self.in_features = in_features
+    ...        self.out_features = out_features
+    ...        self.weight = nn.Parameter(ones(in_features, out_features))
+    ...        self.bias = nn.Parameter(ones(out_features))
+    ...    def forward(self, input):
+    ...        return (input @ self.weight) + self.bias
+    ...
+    >>> config_dict = {
+    ...     "algorithm": "custom",
+    ...     "loss": "mse",
+    ...     "model": {
+    ...         "architecture": "Custom",
+    ...         "name": "linear_model",
+    ...         "in_features": 10,
+    ...         "out_features": 5,
+    ...     }
+    ... }
+    >>> config = AlgorithmModel(**config_dict)
+    """
+
+    # Pydantic class configuration
+    model_config = ConfigDict(
+        protected_namespaces=(),  # allows to use model_* as a field name
+        validate_assignment=True,
+    )
+
+    # Mandatory fields
+    algorithm: Literal["n2v", "care", "n2n", "custom"]  # defined in SupportedAlgorithm
+    loss: Literal["n2v", "mae", "mse"]
+    model: Union[UNetModel, VAEModel, CustomModel] = Field(discriminator="architecture")
+
+    # Optional fields
+    optimizer: OptimizerModel = OptimizerModel()
+    lr_scheduler: LrSchedulerModel = LrSchedulerModel()
+
+    @model_validator(mode="after")
+    def algorithm_cross_validation(self: Self) -> Self:
+        """Validate the algorithm model based on `algorithm`.
+
+        N2V:
+        - loss must be n2v
+        - model must be a `UNetModel`
+
+        Returns
+        -------
+        Self
+            The validated model.
+        """
+        # N2V
+        if self.algorithm == "n2v":
+            # n2v is only compatible with the n2v loss
+            if self.loss != "n2v":
+                raise ValueError(
+                    f"Algorithm {self.algorithm} only supports loss `n2v`."
+                )
+
+            # n2v is only compatible with the UNet model
+            if not isinstance(self.model, UNetModel):
+                raise ValueError(
+                    f"Model for algorithm {self.algorithm} must be a `UNetModel`."
+                )
+
+            # n2v requires the number of input and output channels to be the same
+            if self.model.in_channels != self.model.num_classes:
+                raise ValueError(
+                    "N2V requires the same number of input and output channels. Make "
+                    "sure that `in_channels` and `num_classes` are the same."
+                )
+
+        # N2N
+        if self.algorithm == "n2n":
+            # n2n is only compatible with the UNet model
+            if not isinstance(self.model, UNetModel):
+                raise ValueError(
+                    f"Model for algorithm {self.algorithm} must be a `UNetModel`."
+                )
+
+            # n2n requires the number of input and output channels to be the same
+            if self.model.in_channels != self.model.num_classes:
+                raise ValueError(
+                    "N2N requires the same number of input and output channels. Make "
+                    "sure that `in_channels` and `num_classes` are the same."
+                )
+
+        if self.algorithm == "care" or self.algorithm == "n2n":
+            if self.loss == "n2v":
+                raise ValueError("Supervised algorithms do not support loss `n2v`.")
+
+        if isinstance(self.model, VAEModel):
+            raise ValueError("VAE are currently not implemented.")
+
+        return self
+
+    def __str__(self) -> str:
+        """Pretty string representing the configuration.
+
+        Returns
+        -------
+        str
+            Pretty string.
+        """
+        return pformat(self.model_dump())

careamics/config/architectures/__init__.py

@@ -0,0 +1,17 @@
+"""Deep-learning model configurations."""
+
+__all__ = [
+    "ArchitectureModel",
+    "CustomModel",
+    "UNetModel",
+    "VAEModel",
+    "clear_custom_models",
+    "get_custom_model",
+    "register_model",
+]
+
+from .architecture_model import ArchitectureModel
+from .custom_model import CustomModel
+from .register_model import clear_custom_models, get_custom_model, register_model
+from .unet_model import UNetModel
+from .vae_model import VAEModel

careamics/config/architectures/architecture_model.py

@@ -0,0 +1,29 @@
+from typing import Any, Dict
+
+from pydantic import BaseModel
+
+
+class ArchitectureModel(BaseModel):
+    """
+    Base Pydantic model for all model architectures.
+
+    The `model_dump` method allows removing the `architecture` key from the model.
+    """
+
+    architecture: str
+
+    def model_dump(self, **kwargs: Any) -> Dict[str, Any]:
+        """
+        Dump the model as a dictionary, ignoring the architecture keyword.
+
+        Returns
+        -------
+        dict[str, Any]
+            Model as a dictionnary.
+        """
+        model_dict = super().model_dump(**kwargs)
+
+        # remove the architecture key
+        model_dict.pop("architecture")
+
+        return model_dict

careamics/config/architectures/custom_model.py

@@ -0,0 +1,150 @@
+from __future__ import annotations
+
+from pprint import pformat
+from typing import Any, Dict, Literal
+
+from pydantic import ConfigDict, field_validator, model_validator
+from torch.nn import Module
+from typing_extensions import Self
+
+from .architecture_model import ArchitectureModel
+from .register_model import get_custom_model
+
+
+class CustomModel(ArchitectureModel):
+    """Custom model configuration.
+
+    This Pydantic model allows storing parameters for a custom model. In order for the
+    model to be valid, the specific model needs to be registered using the
+    `register_model` decorator, and its name correctly passed to this model
+    configuration (see Examples).
+
+    Attributes
+    ----------
+    architecture : Literal["Custom"]
+        Discriminator for the custom model, must be set to "Custom".
+    name : str
+        Name of the custom model.
+    parameters : CustomParametersModel
+        Parameters of the custom model.
+
+    Raises
+    ------
+    ValueError
+        If the custom model `name` is unknown.
+    ValueError
+        If the custom model is not a torch Module subclass.
+    ValueError
+        If the custom model parameters are not valid.
+
+    Examples
+    --------
+    >>> from torch import nn, ones
+    >>> from careamics.config import CustomModel, register_model
+    >>> # Register a custom model
+    >>> @register_model(name="my_linear")
+    ... class LinearModel(nn.Module):
+    ...    def __init__(self, in_features, out_features, *args, **kwargs):
+    ...        super().__init__()
+    ...        self.in_features = in_features
+    ...        self.out_features = out_features
+    ...        self.weight = nn.Parameter(ones(in_features, out_features))
+    ...        self.bias = nn.Parameter(ones(out_features))
+    ...    def forward(self, input):
+    ...        return (input @ self.weight) + self.bias
+    ...
+    >>> # Create a configuration
+    >>> config_dict = {
+    ...     "architecture": "Custom",
+    ...     "name": "my_linear",
+    ...     "in_features": 10,
+    ...     "out_features": 5,
+    ... }
+    >>> config = CustomModel(**config_dict)
+    """
+
+    # pydantic model config
+    model_config = ConfigDict(
+        extra="allow",
+    )
+
+    # discriminator used for choosing the pydantic model in Model
+    architecture: Literal["Custom"]
+
+    # name of the custom model
+    name: str
+
+    @field_validator("name")
+    @classmethod
+    def custom_model_is_known(cls, value: str) -> str:
+        """Check whether the custom model is known.
+
+        Parameters
+        ----------
+        value : str
+            Name of the custom model as registered using the `@register_model`
+            decorator.
+        """
+        # delegate error to get_custom_model
+        model = get_custom_model(value)
+
+        # check if it is a torch Module subclass
+        if not issubclass(model, Module):
+            raise ValueError(
+                f'Retrieved class {model} with name "{value}" is not a '
+                f"torch.nn.Module subclass."
+            )
+
+        return value
+
+    @model_validator(mode="after")
+    def check_parameters(self: Self) -> Self:
+        """Validate model by instantiating the model with the parameters.
+
+        Returns
+        -------
+        Self
+            The validated model.
+        """
+        # instantiate model
+        try:
+            get_custom_model(self.name)(**self.model_dump())
+        except Exception as e:
+            raise ValueError(
+                f"error while passing parameters to the model {e}. Verify that all "
+                f"mandatory parameters are provided, and that either the {e} accepts "
+                f"*args and **kwargs in its __init__() method, or that no additional"
+                f"parameter is provided."
+            ) from None
+
+        return self
+
+    def __str__(self) -> str:
+        """Pretty string representing the configuration.
+
+        Returns
+        -------
+        str
+            Pretty string.
+        """
+        return pformat(self.model_dump())
+
+    def model_dump(self, **kwargs: Any) -> Dict[str, Any]:
+        """Dump the model configuration.
+
+        Parameters
+        ----------
+        kwargs : Any
+            Additional keyword arguments from Pydantic BaseModel model_dump method.
+
+        Returns
+        -------
+        Dict[str, Any]
+            Model configuration.
+        """
+        model_dict = super().model_dump()
+
+        # remove the name key
+        model_dict.pop("name")
+
+        return model_dict

careamics/config/architectures/register_model.py

@@ -0,0 +1,101 @@
+from typing import Callable
+
+from torch.nn import Module
+
+CUSTOM_MODELS = {}  # dictionary of custom models {"name": __class__}
+
+
+def register_model(name: str) -> Callable:
+    """Decorator used to register a torch.nn.Module class with a given `name`.
+
+    Parameters
+    ----------
+    name : str
+        Name of the model.
+
+    Returns
+    -------
+    Callable
+        Function allowing to instantiate the wrapped Module class.
+
+    Raises
+    ------
+    ValueError
+        If a model is already registered with that name.
+
+    Examples
+    --------
+    ```python
+    @register_model(name="linear")
+    class LinearModel(nn.Module):
+        def __init__(self, in_features, out_features):
+            super().__init__()
+
+            self.weight = nn.Parameter(ones(in_features, out_features))
+            self.bias = nn.Parameter(ones(out_features))
+
+        def forward(self, input):
+            return (input @ self.weight) + self.bias
+    ```
+    """
+    if name is None or name == "":
+        raise ValueError("Model name cannot be empty.")
+
+    if name in CUSTOM_MODELS:
+        raise ValueError(
+            f"Model {name} already exists. Choose a different name or run "
+            f"`clear_custom_models()` to empty the registry."
+        )
+
+    def add_custom_model(model: Module) -> Module:
+        """Add a custom model to the registry and return it.
+
+        Parameters
+        ----------
+        model : Module
+            Module class to register
+
+        Returns
+        -------
+        Module
+            The registered model.
+        """
+        # add model to the registry
+        CUSTOM_MODELS[name] = model
+
+        return model
+
+    return add_custom_model
+
+
+def get_custom_model(name: str) -> Module:
+    """Get the custom model corresponding to `name` from the registry.
+
+    Parameters
+    ----------
+    name : str
+        Name of the model to retrieve.
+
+    Returns
+    -------
+    Module
+        The requested model.
+
+    Raises
+    ------
+    ValueError
+        If the model is not registered.
+    """
+    if name not in CUSTOM_MODELS:
+        raise ValueError(
+            f"Model {name} is unknown. Have you registered it using "
+            f'@register_model("{name}") as decorator?'
+        )
+
+    return CUSTOM_MODELS[name]
+
+
+def clear_custom_models() -> None:
+    """Clear the custom models registry."""
+    # clear dictionary
+    CUSTOM_MODELS.clear()

careamics/config/architectures/unet_model.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+from typing import Literal
+
+from pydantic import ConfigDict, Field, field_validator
+
+from .architecture_model import ArchitectureModel
+
+
+# TODO tests activation <-> pydantic model, test the literals!
+# TODO annotations for the json schema?
+class UNetModel(ArchitectureModel):
+    """
+    Pydantic model for a N2V(2)-compatible UNet.
+
+    Attributes
+    ----------
+    depth : int
+        Depth of the model, between 1 and 10 (default 2).
+    num_channels_init : int
+        Number of filters of the first level of the network, should be even
+        and minimum 8 (default 96).
+    """
+
+    # pydantic model config
+    model_config = ConfigDict(validate_assignment=True)
+
+    # discriminator used for choosing the pydantic model in Model
+    architecture: Literal["UNet"]
+
+    # parameters
+    # validate_defaults allow ignoring default values in the dump if they were not set
+    conv_dims: Literal[2, 3] = Field(default=2, validate_default=True)
+    num_classes: int = Field(default=1, ge=1, validate_default=True)
+    in_channels: int = Field(default=1, ge=1, validate_default=True)
+    depth: int = Field(default=2, ge=1, le=10, validate_default=True)
+    num_channels_init: int = Field(default=32, ge=8, le=1024, validate_default=True)
+    final_activation: Literal[
+        "None", "Sigmoid", "Softmax", "Tanh", "ReLU", "LeakyReLU"
+    ] = Field(default="None", validate_default=True)
+    n2v2: bool = Field(default=False, validate_default=True)
+
+    @field_validator("num_channels_init")
+    @classmethod
+    def validate_num_channels_init(cls, num_channels_init: int) -> int:
+        """
+        Validate that num_channels_init is even.
+
+        Parameters
+        ----------
+        num_channels_init : int
+            Number of channels.
+
+        Returns
+        -------
+        int
+            Validated number of channels.
+
+        Raises
+        ------
+        ValueError
+            If the number of channels is odd.
+        """
+        # if odd
+        if num_channels_init % 2 != 0:
+            raise ValueError(
+                f"Number of channels for the bottom layer must be even"
+                f" (got {num_channels_init})."
+            )
+
+        return num_channels_init
+
+    def set_3D(self, is_3D: bool) -> None:
+        """
+        Set 3D model by setting the `conv_dims` parameters.
+
+        Parameters
+        ----------
+        is_3D : bool
+            Whether the algorithm is 3D or not.
+        """
+        if is_3D:
+            self.conv_dims = 3
+        else:
+            self.conv_dims = 2
+
+    def is_3D(self) -> bool:
+        """
+        Return whether the model is 3D or not.
+
+        Returns
+        -------
+        bool
+            Whether the model is 3D or not.
+        """
+        return self.conv_dims == 3

careamics/config/architectures/vae_model.py

@@ -0,0 +1,39 @@
+from typing import Literal
+
+from pydantic import (
+    ConfigDict,
+)
+
+from .architecture_model import ArchitectureModel
+
+
+class VAEModel(ArchitectureModel):
+    """VAE model placeholder."""
+
+    model_config = ConfigDict(
+        use_enum_values=True, protected_namespaces=(), validate_assignment=True
+    )
+
+    architecture: Literal["VAE"]
+
+    def set_3D(self, is_3D: bool) -> None:
+        """
+        Set 3D model by setting the `conv_dims` parameters.
+
+        Parameters
+        ----------
+        is_3D : bool
+            Whether the algorithm is 3D or not.
+        """
+        raise NotImplementedError("VAE is not implemented yet.")
+
+    def is_3D(self) -> bool:
+        """
+        Return whether the model is 3D or not.
+
+        Returns
+        -------
+        bool
+            Whether the model is 3D or not.
+        """
+        raise NotImplementedError("VAE is not implemented yet.")