autoden ¶

Auto-Denoise package.

Unsupervised and self-supervised CNN denoising methods.

Modules:

algorithms –

Implementation of various unsupervised and self-supervised denoising methods.
cli –

Module that contains the command line application.
debug –

Debugging utilities.
losses –

Data losses definitions.
models –

Models sub-package.

Classes:

DIP –

Deep image prior.
DataScaleBias –

Data scale and bias.
Denoiser –

Denoising images.
LossRegularizer –

Base class for the regularizer losses.
LossTV –

Total Variation loss function.
N2N –

Self-supervised denoising from pairs of images.
N2V –

Self-supervised denoising from single images.
NetworkParams –

Abstract base class for storing network parameters.

Functions:

compute_scaling_selfsupervised –

Compute input data scaling and bias for self-supervised learning.
compute_scaling_supervised –

Compute input and target data scaling and bias for supervised learning.
create_network –

Create and return a neural network model based on the provided network configuration.
create_optimizer –

Instantiates the desired optimizer for the given model.
fix_invalid_gradient_values –

Fixes invalid gradient values in the model's parameters.
get_num_parameters –

Returns the number of trainable parameters in the model.
load_model_state –

Load a model from disk.
save_model_state –

Save a model's state to disk.

DIP ¶

DIP(
    model: int | str | NetworkParams | Module | Mapping,
    data_scale_bias: DataScaleBias | None = None,
    reg_val: float | LossRegularizer | None = 1e-05,
    device: str = "cuda" if is_available() else "cpu",
    save_epochs_dir: str | None = None,
    verbose: bool = True,
)

Bases: Denoiser

Deep image prior.

Parameters:

model (str | NetworkParams | Module | Mapping | None) –

Type of neural network to use or a specific network (or state) to use
data_scale_bias (DataScaleBias | None, default: None ) –

Scale and bias of the input data, by default None
reg_val (float | None, default: 1e-05 ) –

Regularization value, by default 1e-5
device (str, default: 'cuda' if is_available() else 'cpu' ) –

Device to use, by default "cuda" if cuda is available, otherwise "cpu"
save_epochs_dir (str | None, default: None ) –

Directory where to save network states at each epoch. If None disabled, by default None
verbose (bool, default: True ) –

Whether to produce verbose output, by default True

Methods:

infer –

Inference, given an initial stack of images.
train_supervised –

Supervised training.
train_unsupervised –

Train the model in an unsupervised manner.

Source code in src/autoden/algorithms.py

def __init__(
    self,
    model: int | str | NetworkParams | pt.nn.Module | Mapping,
    data_scale_bias: DataScaleBias | None = None,
    reg_val: float | LossRegularizer | None = 1e-5,
    device: str = "cuda" if pt.cuda.is_available() else "cpu",
    save_epochs_dir: str | None = None,
    verbose: bool = True,
) -> None:
    """Initialize the noise2noise method.

    Parameters
    ----------
    model : str | NetworkParams | pt.nn.Module | Mapping | None
        Type of neural network to use or a specific network (or state) to use
    data_scale_bias : DataScaleBias | None, optional
        Scale and bias of the input data, by default None
    reg_val : float | None, optional
        Regularization value, by default 1e-5
    device : str, optional
        Device to use, by default "cuda" if cuda is available, otherwise "cpu"
    save_epochs_dir : str | None, optional
        Directory where to save network states at each epoch.
        If None disabled, by default None
    verbose : bool, optional
        Whether to produce verbose output, by default True
    """
    if isinstance(model, int):
        if self.save_epochs_dir is None:
            raise ValueError("Directory for saving epochs not specified")

        model = load_model_state(self.save_epochs_dir, epoch_num=model)

    if isinstance(model, (str, NetworkParams, Mapping, pt.nn.Module)):
        self.model = create_network(model, device=device)
    else:
        raise ValueError(f"Invalid model {type(model)}")
    if verbose:
        get_num_parameters(self.model, verbose=True)

    self.data_sb = data_scale_bias

    self.reg_val = reg_val
    self.device = device
    self.save_epochs_dir = save_epochs_dir
    self.verbose = verbose

infer ¶

infer(inp: NDArray) -> NDArray

Inference, given an initial stack of images.

Parameters:

inp (NDArray) –

The input stack of images

Returns:

NDArray –

The denoised stack of images

Source code in src/autoden/algorithms.py

def infer(self, inp: NDArray) -> NDArray:
    """Inference, given an initial stack of images.

    Parameters
    ----------
    inp : NDArray
        The input stack of images

    Returns
    -------
    NDArray
        The denoised stack of images
    """
    # Rescale input
    if self.data_sb is not None:
        inp = inp * self.data_sb.scale_inp - self.data_sb.bias_inp

    inp_t = _single_channel_imgs_to_tensor(inp, device=self.device)

    self.model.eval()
    with pt.inference_mode():
        out_t: pt.Tensor = self.model(inp_t)
        output = out_t.to("cpu").numpy().reshape(inp.shape)

    # Rescale output
    if self.data_sb is not None:
        output = (output + self.data_sb.bias_out) / self.data_sb.scale_out

    return output

train_supervised ¶

train_supervised(
    inp: NDArray,
    tgt: NDArray,
    epochs: int,
    tst_inds: Sequence[int] | NDArray,
    algo: str = "adam",
)

Supervised training.

Parameters:

inp (NDArray) –

The input images
tgt (NDArray) –

The target images
epochs (int) –

Number of training epochs
tst_inds (Sequence[int] | NDArray) –

The validation set indices
algo (str, default: 'adam' ) –

Learning algorithm to use, by default "adam"

Source code in src/autoden/algorithms.py

def train_supervised(
    self,
    inp: NDArray,
    tgt: NDArray,
    epochs: int,
    tst_inds: Sequence[int] | NDArray,
    algo: str = "adam",
):
    """Supervised training.

    Parameters
    ----------
    inp : NDArray
        The input images
    tgt : NDArray
        The target images
    epochs : int
        Number of training epochs
    tst_inds : Sequence[int] | NDArray
        The validation set indices
    algo : str, optional
        Learning algorithm to use, by default "adam"
    """
    num_imgs = inp.shape[0]
    tst_inds = np.array(tst_inds, dtype=int)
    if np.any(tst_inds < 0) or np.any(tst_inds >= num_imgs):
        raise ValueError(
            f"Each cross-validation index should be greater or equal than 0, and less than the number of images {num_imgs}"
        )
    trn_inds = np.delete(np.arange(num_imgs), obj=tst_inds)

    if tgt.ndim == (inp.ndim - 1):
        tgt = np.tile(tgt[None, ...], [num_imgs, *np.ones_like(tgt.shape)])

    if self.data_sb is None:
        self.data_sb = compute_scaling_supervised(inp, tgt)

    # Rescale the datasets
    inp = inp * self.data_sb.scale_inp - self.data_sb.bias_inp
    tgt = tgt * self.data_sb.scale_tgt - self.data_sb.bias_tgt

    # Create datasets
    dset_trn = (inp[trn_inds], tgt[trn_inds])
    dset_tst = (inp[tst_inds], tgt[tst_inds])

    reg = self._get_regularization()
    losses = self._train_selfsimilar(dset_trn, dset_tst, epochs=epochs, algo=algo, regularizer=reg)

    if self.verbose:
        self._plot_loss_curves(losses, f"Supervised {algo.upper()}")

train_unsupervised ¶

train_unsupervised(
    tgt: NDArray,
    epochs: int,
    inp: NDArray | None = None,
    num_tst_ratio: float = 0.2,
    algo: str = "adam",
) -> NDArray

Train the model in an unsupervised manner.

Parameters:

tgt (NDArray) –

The target image to be denoised.
epochs (int) –

The number of training epochs.
inp (NDArray | None, default: None ) –

The input image. If None, a random image will be generated. Default is None.
num_tst_ratio (float, default: 0.2 ) –

The ratio of the test set size to the total dataset size. Default is 0.2.
algo (str, default: 'adam' ) –

The optimization algorithm to use. Default is "adam".

Returns:

NDArray –

The denoised input image.

Notes

This method trains the model using the deep image prior approach in an unsupervised manner. It uses a random initialization for the input image if not provided and applies a scaling and bias transformation to the input and target images. It then splits the data into training and test sets based on the provided ratio and trains the model using the specified optimization algorithm.

Source code in src/autoden/algorithms.py

def train_unsupervised(
    self, tgt: NDArray, epochs: int, inp: NDArray | None = None, num_tst_ratio: float = 0.2, algo: str = "adam"
) -> NDArray:
    """
    Train the model in an unsupervised manner.

    Parameters
    ----------
    tgt : NDArray
        The target image to be denoised.
    epochs : int
        The number of training epochs.
    inp : NDArray | None, optional
        The input image. If None, a random image will be generated.
        Default is None.
    num_tst_ratio : float, optional
        The ratio of the test set size to the total dataset size.
        Default is 0.2.
    algo : str, optional
        The optimization algorithm to use. Default is "adam".

    Returns
    -------
    NDArray
        The denoised input image.

    Notes
    -----
    This method trains the model using the deep image prior approach in an unsupervised manner.
    It uses a random initialization for the input image if not provided and applies a scaling and bias
    transformation to the input and target images. It then splits the data into training and test sets
    based on the provided ratio and trains the model using the specified optimization algorithm.
    """
    if inp is None:
        inp = np.random.normal(size=tgt.shape[-2:], scale=0.25).astype(tgt.dtype)

    if self.data_sb is None:
        self.data_sb = compute_scaling_supervised(inp, tgt)

    # Rescale the datasets
    tmp_inp = inp * self.data_sb.scale_inp - self.data_sb.bias_inp
    tmp_tgt = tgt * self.data_sb.scale_tgt - self.data_sb.bias_tgt

    mask_trn = np.ones_like(tgt, dtype=bool)
    rnd_inds = np.random.random_integers(low=0, high=mask_trn.size - 1, size=int(mask_trn.size * num_tst_ratio))
    mask_trn[np.unravel_index(rnd_inds, shape=mask_trn.shape)] = False

    reg = self._get_regularization()
    losses = self._train_pixelmask_small(tmp_inp, tmp_tgt, mask_trn, epochs=epochs, algo=algo, regularizer=reg)

    if self.verbose:
        self._plot_loss_curves(losses, f"Unsupervised {self.__class__.__name__} {algo.upper()}")

    return inp

DataScaleBias `dataclass` ¶

DataScaleBias(
    scale_inp: float | NDArray = 1.0,
    scale_out: float | NDArray = 1.0,
    scale_tgt: float | NDArray = 1.0,
    bias_inp: float | NDArray = 0.0,
    bias_out: float | NDArray = 0.0,
    bias_tgt: float | NDArray = 0.0,
)

Data scale and bias.

Denoiser ¶

Denoiser(
    model: int | str | NetworkParams | Module | Mapping,
    data_scale_bias: DataScaleBias | None = None,
    reg_val: float | LossRegularizer | None = 1e-05,
    device: str = "cuda" if is_available() else "cpu",
    save_epochs_dir: str | None = None,
    verbose: bool = True,
)

Denoising images.

Parameters:

model (str | NetworkParams | Module | Mapping | None) –

Type of neural network to use or a specific network (or state) to use
data_scale_bias (DataScaleBias | None, default: None ) –

Scale and bias of the input data, by default None
reg_val (float | None, default: 1e-05 ) –

Regularization value, by default 1e-5
device (str, default: 'cuda' if is_available() else 'cpu' ) –

Device to use, by default "cuda" if cuda is available, otherwise "cpu"
save_epochs_dir (str | None, default: None ) –

Directory where to save network states at each epoch. If None disabled, by default None
verbose (bool, default: True ) –

Whether to produce verbose output, by default True

Methods:

infer –

Inference, given an initial stack of images.
train_supervised –

Supervised training.

Source code in src/autoden/algorithms.py

def __init__(
    self,
    model: int | str | NetworkParams | pt.nn.Module | Mapping,
    data_scale_bias: DataScaleBias | None = None,
    reg_val: float | LossRegularizer | None = 1e-5,
    device: str = "cuda" if pt.cuda.is_available() else "cpu",
    save_epochs_dir: str | None = None,
    verbose: bool = True,
) -> None:
    """Initialize the noise2noise method.

    Parameters
    ----------
    model : str | NetworkParams | pt.nn.Module | Mapping | None
        Type of neural network to use or a specific network (or state) to use
    data_scale_bias : DataScaleBias | None, optional
        Scale and bias of the input data, by default None
    reg_val : float | None, optional
        Regularization value, by default 1e-5
    device : str, optional
        Device to use, by default "cuda" if cuda is available, otherwise "cpu"
    save_epochs_dir : str | None, optional
        Directory where to save network states at each epoch.
        If None disabled, by default None
    verbose : bool, optional
        Whether to produce verbose output, by default True
    """
    if isinstance(model, int):
        if self.save_epochs_dir is None:
            raise ValueError("Directory for saving epochs not specified")

        model = load_model_state(self.save_epochs_dir, epoch_num=model)

    if isinstance(model, (str, NetworkParams, Mapping, pt.nn.Module)):
        self.model = create_network(model, device=device)
    else:
        raise ValueError(f"Invalid model {type(model)}")
    if verbose:
        get_num_parameters(self.model, verbose=True)

    self.data_sb = data_scale_bias

    self.reg_val = reg_val
    self.device = device
    self.save_epochs_dir = save_epochs_dir
    self.verbose = verbose

infer ¶

infer(inp: NDArray) -> NDArray

Inference, given an initial stack of images.

Parameters:

inp (NDArray) –

The input stack of images

Returns:

NDArray –

The denoised stack of images

Source code in src/autoden/algorithms.py

def infer(self, inp: NDArray) -> NDArray:
    """Inference, given an initial stack of images.

    Parameters
    ----------
    inp : NDArray
        The input stack of images

    Returns
    -------
    NDArray
        The denoised stack of images
    """
    # Rescale input
    if self.data_sb is not None:
        inp = inp * self.data_sb.scale_inp - self.data_sb.bias_inp

    inp_t = _single_channel_imgs_to_tensor(inp, device=self.device)

    self.model.eval()
    with pt.inference_mode():
        out_t: pt.Tensor = self.model(inp_t)
        output = out_t.to("cpu").numpy().reshape(inp.shape)

    # Rescale output
    if self.data_sb is not None:
        output = (output + self.data_sb.bias_out) / self.data_sb.scale_out

    return output

train_supervised ¶

train_supervised(
    inp: NDArray,
    tgt: NDArray,
    epochs: int,
    tst_inds: Sequence[int] | NDArray,
    algo: str = "adam",
)

Supervised training.

Parameters:

inp (NDArray) –

The input images
tgt (NDArray) –

The target images
epochs (int) –

Number of training epochs
tst_inds (Sequence[int] | NDArray) –

The validation set indices
algo (str, default: 'adam' ) –

Learning algorithm to use, by default "adam"

Source code in src/autoden/algorithms.py

def train_supervised(
    self,
    inp: NDArray,
    tgt: NDArray,
    epochs: int,
    tst_inds: Sequence[int] | NDArray,
    algo: str = "adam",
):
    """Supervised training.

    Parameters
    ----------
    inp : NDArray
        The input images
    tgt : NDArray
        The target images
    epochs : int
        Number of training epochs
    tst_inds : Sequence[int] | NDArray
        The validation set indices
    algo : str, optional
        Learning algorithm to use, by default "adam"
    """
    num_imgs = inp.shape[0]
    tst_inds = np.array(tst_inds, dtype=int)
    if np.any(tst_inds < 0) or np.any(tst_inds >= num_imgs):
        raise ValueError(
            f"Each cross-validation index should be greater or equal than 0, and less than the number of images {num_imgs}"
        )
    trn_inds = np.delete(np.arange(num_imgs), obj=tst_inds)

    if tgt.ndim == (inp.ndim - 1):
        tgt = np.tile(tgt[None, ...], [num_imgs, *np.ones_like(tgt.shape)])

    if self.data_sb is None:
        self.data_sb = compute_scaling_supervised(inp, tgt)

    # Rescale the datasets
    inp = inp * self.data_sb.scale_inp - self.data_sb.bias_inp
    tgt = tgt * self.data_sb.scale_tgt - self.data_sb.bias_tgt

    # Create datasets
    dset_trn = (inp[trn_inds], tgt[trn_inds])
    dset_tst = (inp[tst_inds], tgt[tst_inds])

    reg = self._get_regularization()
    losses = self._train_selfsimilar(dset_trn, dset_tst, epochs=epochs, algo=algo, regularizer=reg)

    if self.verbose:
        self._plot_loss_curves(losses, f"Supervised {algo.upper()}")

LossRegularizer ¶

Bases: MSELoss

Base class for the regularizer losses.

LossTV ¶

LossTV(
    lambda_val: float,
    size_average=None,
    reduce=None,
    reduction: str = "mean",
    isotropic: bool = True,
    ndims: int = 2,
)

Bases: LossRegularizer

Total Variation loss function.

Methods:

forward –

Compute total variation statistics on current batch.

Source code in src/autoden/losses.py

def __init__(
    self,
    lambda_val: float,
    size_average=None,
    reduce=None,
    reduction: str = "mean",
    isotropic: bool = True,
    ndims: int = 2,
) -> None:
    super().__init__(size_average, reduce, reduction)
    self.lambda_val = lambda_val
    self.isotropic = isotropic
    self.ndims = ndims

forward ¶

forward(img: Tensor) -> Tensor

Compute total variation statistics on current batch.

Source code in src/autoden/losses.py

def forward(self, img: pt.Tensor) -> pt.Tensor:
    """Compute total variation statistics on current batch."""
    _check_input_tensor(img, self.ndims)
    axes = list(range(-(self.ndims + 1), 0))

    diffs = [_differentiate(img, dim=dim, position="post") for dim in range(-self.ndims, 0)]
    diffs = pt.stack(diffs, dim=0)

    if self.isotropic:
        # tv_val = pt.sqrt(pt.stack([pt.pow(d, 2) for d in diffs], dim=0).sum(dim=0))
        tv_val = pt.sqrt(pt.pow(diffs, 2).sum(dim=0))
    else:
        # tv_val = pt.stack([d.abs() for d in diffs], dim=0).sum(dim=0)
        tv_val = diffs.abs().sum(dim=0)

    return self.lambda_val * tv_val.sum(axes).mean()

N2N ¶

N2N(
    model: int | str | NetworkParams | Module | Mapping,
    data_scale_bias: DataScaleBias | None = None,
    reg_val: float | LossRegularizer | None = 1e-05,
    device: str = "cuda" if is_available() else "cpu",
    save_epochs_dir: str | None = None,
    verbose: bool = True,
)

Bases: Denoiser

Self-supervised denoising from pairs of images.

Parameters:

model (str | NetworkParams | Module | Mapping | None) –

Type of neural network to use or a specific network (or state) to use
data_scale_bias (DataScaleBias | None, default: None ) –

Scale and bias of the input data, by default None
reg_val (float | None, default: 1e-05 ) –

Regularization value, by default 1e-5
device (str, default: 'cuda' if is_available() else 'cpu' ) –

Device to use, by default "cuda" if cuda is available, otherwise "cpu"
save_epochs_dir (str | None, default: None ) –

Directory where to save network states at each epoch. If None disabled, by default None
verbose (bool, default: True ) –

Whether to produce verbose output, by default True

Methods:

infer –

Inference, given an initial stack of images.
train_selfsupervised –

Train the denoiser using the Noise2Noise self-supervised approach.
train_supervised –

Supervised training.

Source code in src/autoden/algorithms.py

def __init__(
    self,
    model: int | str | NetworkParams | pt.nn.Module | Mapping,
    data_scale_bias: DataScaleBias | None = None,
    reg_val: float | LossRegularizer | None = 1e-5,
    device: str = "cuda" if pt.cuda.is_available() else "cpu",
    save_epochs_dir: str | None = None,
    verbose: bool = True,
) -> None:
    """Initialize the noise2noise method.

    Parameters
    ----------
    model : str | NetworkParams | pt.nn.Module | Mapping | None
        Type of neural network to use or a specific network (or state) to use
    data_scale_bias : DataScaleBias | None, optional
        Scale and bias of the input data, by default None
    reg_val : float | None, optional
        Regularization value, by default 1e-5
    device : str, optional
        Device to use, by default "cuda" if cuda is available, otherwise "cpu"
    save_epochs_dir : str | None, optional
        Directory where to save network states at each epoch.
        If None disabled, by default None
    verbose : bool, optional
        Whether to produce verbose output, by default True
    """
    if isinstance(model, int):
        if self.save_epochs_dir is None:
            raise ValueError("Directory for saving epochs not specified")

        model = load_model_state(self.save_epochs_dir, epoch_num=model)

    if isinstance(model, (str, NetworkParams, Mapping, pt.nn.Module)):
        self.model = create_network(model, device=device)
    else:
        raise ValueError(f"Invalid model {type(model)}")
    if verbose:
        get_num_parameters(self.model, verbose=True)

    self.data_sb = data_scale_bias

    self.reg_val = reg_val
    self.device = device
    self.save_epochs_dir = save_epochs_dir
    self.verbose = verbose

infer ¶

infer(inp: NDArray) -> NDArray

Inference, given an initial stack of images.

Parameters:

inp (NDArray) –

The input stack of images

Returns:

NDArray –

The denoised stack of images

Source code in src/autoden/algorithms.py

def infer(self, inp: NDArray) -> NDArray:
    """Inference, given an initial stack of images.

    Parameters
    ----------
    inp : NDArray
        The input stack of images

    Returns
    -------
    NDArray
        The denoised stack of images
    """
    # Rescale input
    if self.data_sb is not None:
        inp = inp * self.data_sb.scale_inp - self.data_sb.bias_inp

    inp_t = _single_channel_imgs_to_tensor(inp, device=self.device)

    self.model.eval()
    with pt.inference_mode():
        out_t: pt.Tensor = self.model(inp_t)
        output = out_t.to("cpu").numpy().reshape(inp.shape)

    # Rescale output
    if self.data_sb is not None:
        output = (output + self.data_sb.bias_out) / self.data_sb.scale_out

    return output

train_selfsupervised ¶

train_selfsupervised(
    inp: NDArray,
    epochs: int,
    num_tst_ratio: float = 0.2,
    strategy: str = "1:X",
    algo: str = "adam",
    lower_limit: float | NDArray | None = None,
) -> None

Train the denoiser using the Noise2Noise self-supervised approach.

Parameters:

inp (NDArray) –

The input data to be used for training. This should be a NumPy array of shape (N, H, W), where N is the number of samples, and H and W are the height and width of each sample, respectively.
epochs (int) –

The number of epochs to train the model.
num_tst_ratio (float, default: 0.2 ) –

The ratio of the input data to be used for testing. The remaining data will be used for training. Default is 0.2.
strategy (str, default: '1:X' ) –

The strategy to be used for creating input-target pairs. The available strategies are: - "1:X": Use the mean of the remaining samples as the target for each sample. - "X:1": Use the mean of the remaining samples as the input for each sample. Default is "1:X".
algo (str, default: 'adam' ) –

The optimization algorithm to be used for training. The available algorithms are: - "adam": Adam optimizer. Default is "adam".
lower_limit (float | NDArray | None, default: None ) –

The lower limit for the input data. If provided, the input data will be clipped to this limit. Default is None.

Notes

This method uses the Noise2Noise self-supervised approach to train the denoiser. The input data is used to generate target data based on the specified strategy. The training process involves creating pairs of input and target data and then training the model to minimize the difference between the predicted and target data.

Source code in src/autoden/algorithms.py

def train_selfsupervised(
    self,
    inp: NDArray,
    epochs: int,
    num_tst_ratio: float = 0.2,
    strategy: str = "1:X",
    algo: str = "adam",
    lower_limit: float | NDArray | None = None,
) -> None:
    """
    Train the denoiser using the Noise2Noise self-supervised approach.

    Parameters
    ----------
    inp : NDArray
        The input data to be used for training. This should be a NumPy array of shape (N, H, W), where N is the
        number of samples, and H and W are the height and width of each sample, respectively.
    epochs : int
        The number of epochs to train the model.
    num_tst_ratio : float, optional
        The ratio of the input data to be used for testing. The remaining data will be used for training.
        Default is 0.2.
    strategy : str, optional
        The strategy to be used for creating input-target pairs. The available strategies are:
        - "1:X": Use the mean of the remaining samples as the target for each sample.
        - "X:1": Use the mean of the remaining samples as the input for each sample.
        Default is "1:X".
    algo : str, optional
        The optimization algorithm to be used for training. The available algorithms are:
        - "adam": Adam optimizer.
        Default is "adam".
    lower_limit : float | NDArray | None, optional
        The lower limit for the input data. If provided, the input data will be clipped to this limit.
        Default is None.

    Notes
    -----
    This method uses the Noise2Noise self-supervised approach to train the denoiser. The input data is used to
    generate target data based on the specified strategy. The training process involves creating pairs of input
    and target data and then training the model to minimize the difference between the predicted and target data.
    """
    if self.data_sb is None:
        self.data_sb = compute_scaling_selfsupervised(inp)

    # Rescale the datasets
    inp = inp * self.data_sb.scale_inp - self.data_sb.bias_inp

    mask_trn = np.ones_like(inp, dtype=bool)
    rnd_inds = np.random.random_integers(low=0, high=mask_trn.size - 1, size=int(mask_trn.size * num_tst_ratio))
    mask_trn[np.unravel_index(rnd_inds, shape=mask_trn.shape)] = False

    inp_x = np.stack([np.delete(inp, obj=ii, axis=0).mean(axis=0) for ii in range(len(inp))], axis=0)
    if strategy.upper() == "1:X":
        tmp_inp = inp
        tmp_tgt = inp_x
    elif strategy.upper() == "X:1":
        tmp_inp = inp_x
        tmp_tgt = inp
    else:
        raise ValueError(f"Strategy {strategy} not implemented. Please choose one of: ['1:X', 'X:1']")

    tmp_inp = tmp_inp.astype(np.float32)
    tmp_tgt = tmp_tgt.astype(np.float32)

    reg = self._get_regularization()
    losses = self._train_pixelmask_small(
        tmp_inp, tmp_tgt, mask_trn, epochs=epochs, algo=algo, regularizer=reg, lower_limit=lower_limit
    )

    if self.verbose:
        self._plot_loss_curves(losses, f"Self-supervised {self.__class__.__name__} {algo.upper()}")

train_supervised ¶

train_supervised(
    inp: NDArray,
    tgt: NDArray,
    epochs: int,
    tst_inds: Sequence[int] | NDArray,
    algo: str = "adam",
)

Supervised training.

Parameters:

inp (NDArray) –

The input images
tgt (NDArray) –

The target images
epochs (int) –

Number of training epochs
tst_inds (Sequence[int] | NDArray) –

The validation set indices
algo (str, default: 'adam' ) –

Learning algorithm to use, by default "adam"

Source code in src/autoden/algorithms.py

def train_supervised(
    self,
    inp: NDArray,
    tgt: NDArray,
    epochs: int,
    tst_inds: Sequence[int] | NDArray,
    algo: str = "adam",
):
    """Supervised training.

    Parameters
    ----------
    inp : NDArray
        The input images
    tgt : NDArray
        The target images
    epochs : int
        Number of training epochs
    tst_inds : Sequence[int] | NDArray
        The validation set indices
    algo : str, optional
        Learning algorithm to use, by default "adam"
    """
    num_imgs = inp.shape[0]
    tst_inds = np.array(tst_inds, dtype=int)
    if np.any(tst_inds < 0) or np.any(tst_inds >= num_imgs):
        raise ValueError(
            f"Each cross-validation index should be greater or equal than 0, and less than the number of images {num_imgs}"
        )
    trn_inds = np.delete(np.arange(num_imgs), obj=tst_inds)

    if tgt.ndim == (inp.ndim - 1):
        tgt = np.tile(tgt[None, ...], [num_imgs, *np.ones_like(tgt.shape)])

    if self.data_sb is None:
        self.data_sb = compute_scaling_supervised(inp, tgt)

    # Rescale the datasets
    inp = inp * self.data_sb.scale_inp - self.data_sb.bias_inp
    tgt = tgt * self.data_sb.scale_tgt - self.data_sb.bias_tgt

    # Create datasets
    dset_trn = (inp[trn_inds], tgt[trn_inds])
    dset_tst = (inp[tst_inds], tgt[tst_inds])

    reg = self._get_regularization()
    losses = self._train_selfsimilar(dset_trn, dset_tst, epochs=epochs, algo=algo, regularizer=reg)

    if self.verbose:
        self._plot_loss_curves(losses, f"Supervised {algo.upper()}")

N2V ¶

N2V(
    model: int | str | NetworkParams | Module | Mapping,
    data_scale_bias: DataScaleBias | None = None,
    reg_val: float | LossRegularizer | None = 1e-05,
    device: str = "cuda" if is_available() else "cpu",
    save_epochs_dir: str | None = None,
    verbose: bool = True,
)

Bases: Denoiser

Self-supervised denoising from single images.

Parameters:

model (str | NetworkParams | Module | Mapping | None) –

Type of neural network to use or a specific network (or state) to use
data_scale_bias (DataScaleBias | None, default: None ) –

Scale and bias of the input data, by default None
reg_val (float | None, default: 1e-05 ) –

Regularization value, by default 1e-5
device (str, default: 'cuda' if is_available() else 'cpu' ) –

Device to use, by default "cuda" if cuda is available, otherwise "cpu"
save_epochs_dir (str | None, default: None ) –

Directory where to save network states at each epoch. If None disabled, by default None
verbose (bool, default: True ) –

Whether to produce verbose output, by default True

Methods:

infer –

Inference, given an initial stack of images.
train_selfsupervised –

Self-supervised training.
train_supervised –

Supervised training.

Source code in src/autoden/algorithms.py

def __init__(
    self,
    model: int | str | NetworkParams | pt.nn.Module | Mapping,
    data_scale_bias: DataScaleBias | None = None,
    reg_val: float | LossRegularizer | None = 1e-5,
    device: str = "cuda" if pt.cuda.is_available() else "cpu",
    save_epochs_dir: str | None = None,
    verbose: bool = True,
) -> None:
    """Initialize the noise2noise method.

    Parameters
    ----------
    model : str | NetworkParams | pt.nn.Module | Mapping | None
        Type of neural network to use or a specific network (or state) to use
    data_scale_bias : DataScaleBias | None, optional
        Scale and bias of the input data, by default None
    reg_val : float | None, optional
        Regularization value, by default 1e-5
    device : str, optional
        Device to use, by default "cuda" if cuda is available, otherwise "cpu"
    save_epochs_dir : str | None, optional
        Directory where to save network states at each epoch.
        If None disabled, by default None
    verbose : bool, optional
        Whether to produce verbose output, by default True
    """
    if isinstance(model, int):
        if self.save_epochs_dir is None:
            raise ValueError("Directory for saving epochs not specified")

        model = load_model_state(self.save_epochs_dir, epoch_num=model)

    if isinstance(model, (str, NetworkParams, Mapping, pt.nn.Module)):
        self.model = create_network(model, device=device)
    else:
        raise ValueError(f"Invalid model {type(model)}")
    if verbose:
        get_num_parameters(self.model, verbose=True)

    self.data_sb = data_scale_bias

    self.reg_val = reg_val
    self.device = device
    self.save_epochs_dir = save_epochs_dir
    self.verbose = verbose

infer ¶

infer(inp: NDArray) -> NDArray

Inference, given an initial stack of images.

Parameters:

inp (NDArray) –

The input stack of images

Returns:

NDArray –

The denoised stack of images

Source code in src/autoden/algorithms.py

def infer(self, inp: NDArray) -> NDArray:
    """Inference, given an initial stack of images.

    Parameters
    ----------
    inp : NDArray
        The input stack of images

    Returns
    -------
    NDArray
        The denoised stack of images
    """
    # Rescale input
    if self.data_sb is not None:
        inp = inp * self.data_sb.scale_inp - self.data_sb.bias_inp

    inp_t = _single_channel_imgs_to_tensor(inp, device=self.device)

    self.model.eval()
    with pt.inference_mode():
        out_t: pt.Tensor = self.model(inp_t)
        output = out_t.to("cpu").numpy().reshape(inp.shape)

    # Rescale output
    if self.data_sb is not None:
        output = (output + self.data_sb.bias_out) / self.data_sb.scale_out

    return output

train_selfsupervised ¶

train_selfsupervised(
    inp: NDArray,
    epochs: int,
    tst_inds: Sequence[int] | NDArray,
    mask_shape: int | Sequence[int] | NDArray = 1,
    ratio_blind_spot: float = 0.015,
    algo: str = "adam",
)

Self-supervised training.

Parameters:

inp (NDArray) –

The input images, which will also be targets
epochs (int) –

Number of training epochs
tst_inds (Sequence[int] | NDArray) –

The validation set indices
mask_shape (int | Sequence[int] | NDArray, default: 1 ) –

Shape of the blind spot mask, by default 1.
algo (str, default: 'adam' ) –

Learning algorithm to use, by default "adam"

Source code in src/autoden/algorithms.py

def train_selfsupervised(
    self,
    inp: NDArray,
    epochs: int,
    tst_inds: Sequence[int] | NDArray,
    mask_shape: int | Sequence[int] | NDArray = 1,
    ratio_blind_spot: float = 0.015,
    algo: str = "adam",
):
    """Self-supervised training.

    Parameters
    ----------
    inp : NDArray
        The input images, which will also be targets
    epochs : int
        Number of training epochs
    tst_inds : Sequence[int] | NDArray
        The validation set indices
    mask_shape : int | Sequence[int] | NDArray
        Shape of the blind spot mask, by default 1.
    algo : str, optional
        Learning algorithm to use, by default "adam"
    """
    num_imgs = inp.shape[0]
    tst_inds = np.array(tst_inds, dtype=int)
    if np.any(tst_inds < 0) or np.any(tst_inds >= num_imgs):
        raise ValueError(
            f"Each cross-validation index should be greater or equal than 0, and less than the number of images {num_imgs}"
        )
    trn_inds = np.delete(np.arange(num_imgs), obj=tst_inds)

    if self.data_sb is None:
        self.data_sb = compute_scaling_selfsupervised(inp)

    # Rescale the datasets
    inp = inp * self.data_sb.scale_inp - self.data_sb.bias_inp

    inp_trn = inp[trn_inds]
    inp_tst = inp[tst_inds]

    reg = self._get_regularization()
    losses = self._train_n2v_pixelmask_small(
        inp_trn,
        inp_tst,
        epochs=epochs,
        mask_shape=mask_shape,
        ratio_blind_spot=ratio_blind_spot,
        algo=algo,
        regularizer=reg,
    )

    self._plot_loss_curves(losses, f"Self-supervised {self.__class__.__name__} {algo.upper()}")

train_supervised ¶

train_supervised(
    inp: NDArray,
    tgt: NDArray,
    epochs: int,
    tst_inds: Sequence[int] | NDArray,
    algo: str = "adam",
)

Supervised training.

Parameters:

inp (NDArray) –

The input images
tgt (NDArray) –

The target images
epochs (int) –

Number of training epochs
tst_inds (Sequence[int] | NDArray) –

The validation set indices
algo (str, default: 'adam' ) –

Learning algorithm to use, by default "adam"

Source code in src/autoden/algorithms.py

def train_supervised(
    self,
    inp: NDArray,
    tgt: NDArray,
    epochs: int,
    tst_inds: Sequence[int] | NDArray,
    algo: str = "adam",
):
    """Supervised training.

    Parameters
    ----------
    inp : NDArray
        The input images
    tgt : NDArray
        The target images
    epochs : int
        Number of training epochs
    tst_inds : Sequence[int] | NDArray
        The validation set indices
    algo : str, optional
        Learning algorithm to use, by default "adam"
    """
    num_imgs = inp.shape[0]
    tst_inds = np.array(tst_inds, dtype=int)
    if np.any(tst_inds < 0) or np.any(tst_inds >= num_imgs):
        raise ValueError(
            f"Each cross-validation index should be greater or equal than 0, and less than the number of images {num_imgs}"
        )
    trn_inds = np.delete(np.arange(num_imgs), obj=tst_inds)

    if tgt.ndim == (inp.ndim - 1):
        tgt = np.tile(tgt[None, ...], [num_imgs, *np.ones_like(tgt.shape)])

    if self.data_sb is None:
        self.data_sb = compute_scaling_supervised(inp, tgt)

    # Rescale the datasets
    inp = inp * self.data_sb.scale_inp - self.data_sb.bias_inp
    tgt = tgt * self.data_sb.scale_tgt - self.data_sb.bias_tgt

    # Create datasets
    dset_trn = (inp[trn_inds], tgt[trn_inds])
    dset_tst = (inp[tst_inds], tgt[tst_inds])

    reg = self._get_regularization()
    losses = self._train_selfsimilar(dset_trn, dset_tst, epochs=epochs, algo=algo, regularizer=reg)

    if self.verbose:
        self._plot_loss_curves(losses, f"Supervised {algo.upper()}")

NetworkParams ¶

NetworkParams(
    n_features: int,
    n_channels_in: int = 1,
    n_channels_out: int = 1,
)

Bases: ABC

Abstract base class for storing network parameters.

Methods:

get_model –

Get the associated model with the selected parameters.

Source code in src/autoden/models/config.py

def __init__(self, n_features: int, n_channels_in: int = 1, n_channels_out: int = 1) -> None:
    self.n_channels_in = n_channels_in
    self.n_channels_out = n_channels_out
    self.n_features = n_features

get_model `abstractmethod` ¶

get_model(
    device: str = "cuda" if is_available() else "cpu",
) -> Module

Get the associated model with the selected parameters.

Parameters:

device (str, default: 'cuda' if is_available() else 'cpu' ) –

The device that the the model should run on, by default "cuda" if cuda is available, otherwise "cpu".

Returns:

Module –

The model.

Source code in src/autoden/models/config.py

@abstractmethod
def get_model(self, device: str = "cuda" if is_cuda_available() else "cpu") -> Module:
    """Get the associated model with the selected parameters.

    Parameters
    ----------
    device : str, optional
        The device that the the model should run on, by default "cuda" if cuda is available, otherwise "cpu".

    Returns
    -------
    Module
        The model.
    """

compute_scaling_selfsupervised ¶

compute_scaling_selfsupervised(
    inp: NDArray,
) -> DataScaleBias

Compute input data scaling and bias for self-supervised learning.

Parameters:

inp (NDArray) –

Input data.

Returns:

DataScaleBias –

An instance of DataScaleBias containing the computed scaling and bias values.

Source code in src/autoden/algorithms.py

def compute_scaling_selfsupervised(inp: NDArray) -> DataScaleBias:
    """
    Compute input data scaling and bias for self-supervised learning.

    Parameters
    ----------
    inp : NDArray
        Input data.

    Returns
    -------
    DataScaleBias
        An instance of DataScaleBias containing the computed scaling and bias values.
    """
    range_vals_inp = _get_normalization(inp, percentile=0.001)

    sb = DataScaleBias()
    sb.scale_inp = 1 / (range_vals_inp[1] - range_vals_inp[0])
    sb.scale_out = sb.scale_tgt = sb.scale_inp

    sb.bias_inp = range_vals_inp[2] * sb.scale_inp
    sb.bias_out = sb.bias_tgt = sb.bias_inp

    return sb

compute_scaling_supervised ¶

compute_scaling_supervised(
    inp: NDArray, tgt: NDArray
) -> DataScaleBias

Compute input and target data scaling and bias for supervised learning.

Parameters:

inp (NDArray) –

Input data.
tgt (NDArray) –

Target data.

Returns:

DataScaleBias –

An instance of DataScaleBias containing the computed scaling and bias values.

Source code in src/autoden/algorithms.py

def compute_scaling_supervised(inp: NDArray, tgt: NDArray) -> DataScaleBias:
    """
    Compute input and target data scaling and bias for supervised learning.

    Parameters
    ----------
    inp : NDArray
        Input data.
    tgt : NDArray
        Target data.

    Returns
    -------
    DataScaleBias
        An instance of DataScaleBias containing the computed scaling and bias values.
    """
    range_vals_inp = _get_normalization(inp, percentile=0.001)
    range_vals_tgt = _get_normalization(tgt, percentile=0.001)

    sb = DataScaleBias()
    sb.scale_inp = 1 / (range_vals_inp[1] - range_vals_inp[0])
    sb.scale_tgt = 1 / (range_vals_tgt[1] - range_vals_tgt[0])
    sb.scale_out = sb.scale_tgt

    sb.bias_inp = range_vals_inp[2] * sb.scale_inp
    sb.bias_tgt = range_vals_tgt[2] * sb.scale_tgt
    sb.bias_out = sb.bias_tgt

    return sb

create_network ¶

create_network(
    model: str | NetworkParams | Mapping | Module,
    init_params: Mapping | None = None,
    state_dict: Mapping | None = None,
    device: str = "cuda" if is_available() else "cpu",
) -> Module

Create and return a neural network model based on the provided network configuration.

Parameters:

model (str | NetworkParams | Mapping | Module) –

The network configuration. It can be a string specifying the network type, an instance of NetworkParams, or an already instantiated Module. If a string is provided, it must be one of the supported network types: "msd", "unet", or "dncnn".
state_dict (Mapping | None, default: None ) –

A dictionary containing the state dictionary of the model. If provided, the model's parameters will be loaded from this dictionary. Default is None.
device (str, default: 'cuda' if is_available() else 'cpu' ) –

The device to which the model should be moved. Default is "cuda" if CUDA is available, otherwise "cpu".

Returns:

Module –

The created neural network model.

Raises:

ValueError –

If the provided network name is invalid or the network type is not supported.

Notes

The function supports the following network types: - "msd": Multi-Scale Dense Network. - "unet": U-Net. - "dncnn": Denoising Convolutional Neural Network.

Examples:

>>> net = create_network("unet")
>>> print(net)
Model UNet - num. parameters: 1234567

Source code in src/autoden/models/config.py

def create_network(
    model: str | NetworkParams | Mapping | Module,
    init_params: Mapping | None = None,
    state_dict: Mapping | None = None,
    device: str = "cuda" if is_cuda_available() else "cpu",
) -> Module:
    """
    Create and return a neural network model based on the provided network configuration.

    Parameters
    ----------
    model : str | NetworkParams | Mapping | Module
        The network configuration. It can be a string specifying the network type,
        an instance of `NetworkParams`, or an already instantiated `Module`.
        If a string is provided, it must be one of the supported network types:
        "msd", "unet", or "dncnn".
    state_dict : Mapping | None, optional
        A dictionary containing the state dictionary of the model. If provided,
        the model's parameters will be loaded from this dictionary. Default is None.
    device : str, optional
        The device to which the model should be moved. Default is "cuda" if CUDA is available,
        otherwise "cpu".

    Returns
    -------
    Module
        The created neural network model.

    Raises
    ------
    ValueError
        If the provided network name is invalid or the network type is not supported.

    Notes
    -----
    The function supports the following network types:
    - "msd": Multi-Scale Dense Network.
    - "unet": U-Net.
    - "dncnn": Denoising Convolutional Neural Network.

    Examples
    --------
    >>> net = create_network("unet")
    >>> print(net)
    Model UNet - num. parameters: 1234567
    """
    if isinstance(model, Mapping):
        if not all(key in model for key in ("model_class", "init_params", "state_dict")):
            raise ValueError(
                "Malformed model state dictionary. Expected mandatory fields: 'model_class', 'init_params', and 'state_dict'"
            )
        state_dict = model["state_dict"]
        init_params = model["init_params"]
        model = model["model_class"]

    if init_params is None:
        init_params = dict()
    else:
        init_params = dict(**init_params)

    for par in ("device", "verbose"):
        if par in init_params:
            del init_params[par]

    if isinstance(model, str):
        if model.lower() in ("msd", MSDnet.__name__.lower()):
            model = NetworkParamsMSD(**init_params)
        elif model.lower() == UNet.__name__.lower():
            model = NetworkParamsUNet(**init_params)
        elif model.lower() == DnCNN.__name__.lower():
            model = NetworkParamsDnCNN(**init_params)
        elif model.lower() == Resnet.__name__.lower():
            model = NetworkParamsResnet(**init_params)
        else:
            raise ValueError(f"Invalid model name: {model}")

    if isinstance(model, NetworkParams):
        net = model.get_model(device)
    elif isinstance(model, Module):
        net = model.to(device=device)
    else:
        raise ValueError(f"Invalid model type: {type(model)}")

    if state_dict is not None:
        net.load_state_dict(state_dict)
        net.to(device)  # Needed to ensure that the model lives in the correct device

    print(f"Model {net.__class__.__name__} - num. parameters: {sum(p.numel() for p in net.parameters() if p.requires_grad)}")
    return net

create_optimizer ¶

create_optimizer(
    network: Module,
    algo: str = "adam",
    learning_rate: float = 0.001,
    weight_decay: float = 0.01,
    optim_state: Mapping | None = None,
) -> Optimizer

Instantiates the desired optimizer for the given model.

Parameters:

network (Module) –

The network to train.
algo (str, default: 'adam' ) –

The requested optimizer, by default "adam".
learning_rate (float, default: 0.001 ) –

The desired learning rate, by default 1e-3.
weight_decay (float, default: 0.01 ) –

The desired weight decay, by default 1e-2.
optim_state (Mapping | None, default: None ) –

The state dictionary for the optimizer, by default None.

Returns:

Optimizer –

The chosen optimizer.

Raises:

ValueError –

If an unsupported algorithm is requested.

Source code in src/autoden/models/config.py

def create_optimizer(
    network: Module,
    algo: str = "adam",
    learning_rate: float = 1e-3,
    weight_decay: float = 1e-2,
    optim_state: Mapping | None = None,
) -> Optimizer:
    """Instantiates the desired optimizer for the given model.

    Parameters
    ----------
    network : torch.nn.Module
        The network to train.
    algo : str, optional
        The requested optimizer, by default "adam".
    learning_rate : float, optional
        The desired learning rate, by default 1e-3.
    weight_decay : float, optional
        The desired weight decay, by default 1e-2.
    optim_state : Mapping | None, optional
        The state dictionary for the optimizer, by default None.

    Returns
    -------
    torch.optim.Optimizer
        The chosen optimizer.

    Raises
    ------
    ValueError
        If an unsupported algorithm is requested.
    """
    if algo.lower() == "adam":
        optimizer = pt.optim.AdamW(network.parameters(), lr=learning_rate, weight_decay=weight_decay)
    elif algo.lower() == "sgd":
        optimizer = pt.optim.SGD(network.parameters(), lr=learning_rate, weight_decay=weight_decay)
    elif algo.lower() == "rmsprop":
        optimizer = pt.optim.RMSprop(network.parameters(), lr=learning_rate, weight_decay=weight_decay)
    elif algo.lower() == "lbfgs":
        optimizer = pt.optim.LBFGS(network.parameters(), lr=learning_rate, max_iter=10000, history_size=50)
    else:
        raise ValueError(f"Unknown algorithm: {algo}")

    if optim_state is not None:
        optimizer.load_state_dict(dict(**optim_state))

    return optimizer

fix_invalid_gradient_values ¶

fix_invalid_gradient_values(model: Module) -> None

Fixes invalid gradient values in the model's parameters.

This function iterates over all parameters of the given model and sets the gradient values to zero where they are not finite (i.e., NaN or infinity).

Parameters:

model (Module) –

The neural network model whose gradient values need to be fixed.

Returns:

None –

This function modifies the gradients in place and does not return anything.

Source code in src/autoden/models/param_utils.py

def fix_invalid_gradient_values(model: nn.Module) -> None:
    """
    Fixes invalid gradient values in the model's parameters.

    This function iterates over all parameters of the given model and sets the
    gradient values to zero where they are not finite (i.e., NaN or infinity).

    Parameters
    ----------
    model : nn.Module
        The neural network model whose gradient values need to be fixed.

    Returns
    -------
    None
        This function modifies the gradients in place and does not return anything.
    """
    for pars in model.parameters():
        if pars.grad is not None:
            pars.grad[pt.logical_not(pt.isfinite(pars.grad))] = 0.0

get_num_parameters ¶

get_num_parameters(
    model: Module, verbose: bool = False
) -> int

Returns the number of trainable parameters in the model.

Parameters:

model (Module) –

The model to count the parameters for.
verbose (bool, default: False ) –

If True, prints the number of parameters, by default False.

Returns:

int –

The number of trainable parameters.

Source code in src/autoden/models/param_utils.py

def get_num_parameters(model: nn.Module, verbose: bool = False) -> int:
    """Returns the number of trainable parameters in the model.

    Parameters
    ----------
    model : torch.nn.Module
        The model to count the parameters for.
    verbose : bool, optional
        If True, prints the number of parameters, by default False.

    Returns
    -------
    int
        The number of trainable parameters.
    """
    num_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    if verbose:
        print(f"Model {model.__class__.__name__} - num. parameters: {num_params}")
    return num_params

load_model_state ¶

load_model_state(
    save_epochs_dir: str | Path,
    epoch_num: int | None = None,
) -> Mapping

Load a model from disk.

Parameters:

save_epochs_dir (str | Path) –

The director where the models are saved
epoch_num (int | None, default: None ) –

The epoch number or if None/-1 the best state will be loaded, by default None

Returns:

Mapping –

The loaded model state and possibly an optimizer state.

Raises:

ValueError –

When the directory does not exist or the requested model is not available.

Source code in src/autoden/models/io.py

def load_model_state(save_epochs_dir: str | Path, epoch_num: int | None = None) -> Mapping:
    """Load a model from disk.

    Parameters
    ----------
    save_epochs_dir : str | Path
        The director where the models are saved
    epoch_num : int | None, optional
        The epoch number or if None/-1 the best state will be loaded, by default None

    Returns
    -------
    Mapping
        The loaded model state and possibly an optimizer state.

    Raises
    ------
    ValueError
        When the directory does not exist or the requested model is not available.
    """
    epochs_base_path = Path(save_epochs_dir) / "weights"
    if not epochs_base_path.exists():
        raise ValueError(f"Directory of the model state {epochs_base_path} does not exist!")

    if epoch_num is None or epoch_num == -1:
        state_path = epochs_base_path / "weights.pt"
    else:
        state_path = epochs_base_path / f"weights_epoch_{epoch_num}.pt"
    if not state_path.exists():
        raise ValueError(f"Model state {state_path} does not exist!")

    print(f"Loading state path: {state_path}")
    return pt.load(state_path)

save_model_state ¶

save_model_state(
    save_epochs_dir: str | Path,
    epoch_num: int,
    model: Module,
    optim_state: Mapping | None = None,
    is_best: bool = False,
) -> None

Save a model's state to disk.

This function saves the state of a model and optionally its optimizer to disk. The model state is saved in a directory specified by save_epochs_dir. If is_best is True, the model state is saved as "weights.pt". Otherwise, it is saved with a filename that includes the epoch number.

Parameters:

save_epochs_dir (str | Path) –

The directory where to save the model state.
epoch_num (int) –

The epoch number.
model (Module) –

The model whose state is to be saved.
optim_state (Mapping, default: None ) –

The optimizer state to save, by default None.
is_best (bool, default: False ) –

Whether it is the best fitted model, by default False.

Returns:

None –

Source code in src/autoden/models/io.py

def save_model_state(
    save_epochs_dir: str | Path,
    epoch_num: int,
    model: Module,
    optim_state: Mapping | None = None,
    is_best: bool = False,
) -> None:
    """Save a model's state to disk.

    This function saves the state of a model and optionally its optimizer to disk.
    The model state is saved in a directory specified by `save_epochs_dir`. If
    `is_best` is True, the model state is saved as "weights.pt". Otherwise, it is
    saved with a filename that includes the epoch number.

    Parameters
    ----------
    save_epochs_dir : str | Path
        The directory where to save the model state.
    epoch_num : int
        The epoch number.
    model : Module
        The model whose state is to be saved.
    optim_state : Mapping, optional
        The optimizer state to save, by default None.
    is_best : bool, optional
        Whether it is the best fitted model, by default False.

    Returns
    -------
    None
    """
    epochs_base_path = Path(save_epochs_dir) / "weights"
    epochs_base_path.mkdir(parents=True, exist_ok=True)

    dst_file = epochs_base_path / ("weights.pt" if is_best else f"weights_epoch_{epoch_num}.pt")
    save_model(dst_file=dst_file, model=model, optim_state=optim_state, epoch_num=epoch_num)

autoden ¶

DIP ¶

infer ¶

train_supervised ¶

train_unsupervised ¶

DataScaleBias dataclass ¶

Denoiser ¶

infer ¶

train_supervised ¶

LossRegularizer ¶

LossTV ¶

forward ¶

N2N ¶

infer ¶

train_selfsupervised ¶

train_supervised ¶

N2V ¶

infer ¶

train_selfsupervised ¶

train_supervised ¶

NetworkParams ¶

get_model abstractmethod ¶

compute_scaling_selfsupervised ¶

compute_scaling_supervised ¶

create_network ¶

create_optimizer ¶

fix_invalid_gradient_values ¶

get_num_parameters ¶

load_model_state ¶

save_model_state ¶

DataScaleBias `dataclass` ¶

get_model `abstractmethod` ¶