Skip to content

Lightning

Source

Deprecated Lightning modules from CAREamics v0.1.0.

FCNModule

Bases: LightningModule

CAREamics Lightning module.

This class encapsulates the PyTorch model along with the training, validation, and testing logic. It is configured using an AlgorithmModel Pydantic class.

Parameters:

  • algorithm_config (AlgorithmModel or dict) –

    Algorithm configuration.

Attributes:

  • model (Module) –

    PyTorch model.

  • loss_func (Module) –

    Loss function.

  • optimizer_name (str) –

    Optimizer name.

  • optimizer_params (dict) –

    Optimizer parameters.

  • lr_scheduler_name (str) –

    Learning rate scheduler name.

__init__(algorithm_config)

Lightning module for CAREamics.

This class encapsulates the a PyTorch model along with the training, validation, and testing logic. It is configured using an AlgorithmModel Pydantic class.

Parameters:

  • algorithm_config (AlgorithmModel or dict) –

    Algorithm configuration.

configure_optimizers()

Configure optimizers and learning rate schedulers.

Returns:

  • Any

    Optimizer and learning rate scheduler.

forward(x)

Forward pass.

Parameters:

  • x (Any) –

    Input tensor.

Returns:

  • Any

    Output tensor.

predict_step(batch, batch_idx)

Prediction step.

Parameters:

  • batch (Tensor) –

    Input batch.

  • batch_idx (Any) –

    Batch index.

Returns:

  • Any

    Model output.

training_step(batch, batch_idx)

Training step.

Parameters:

  • batch (Tensor) –

    Input batch.

  • batch_idx (Any) –

    Batch index.

Returns:

  • Any

    Loss value.

validation_step(batch, batch_idx)

Validation step.

Parameters:

  • batch (Tensor) –

    Input batch.

  • batch_idx (Any) –

    Batch index.

PredictDataModule

Bases: LightningDataModule

CAREamics Lightning prediction data module.

The data module can be used with Path, str or numpy arrays. The data can be either a folder containing images or a single file.

To read custom data types, you can set data_type to custom in data_config and provide a function that returns a numpy array from a path as read_source_func parameter. The function will receive a Path object and an axies string as arguments, the axes being derived from the data_config.

You can also provide a fnmatch and Path.rglob compatible expression (e.g. "*.czi") to filter the files extension using extension_filter.

Parameters:

  • pred_config (InferenceModel) –

    Pydantic model for CAREamics prediction configuration.

  • pred_data (Path or str or ndarray) –

    Prediction data, can be a path to a folder, a file or a numpy array.

  • read_source_func (Callable, default: None ) –

    Function to read custom types, by default None.

  • extension_filter (str, default: '' ) –

    Filter to filter file extensions for custom types, by default "".

  • dataloader_params (dict, default: None ) –

    Dataloader parameters, by default {}.

__init__(pred_config, pred_data, read_source_func=None, extension_filter='', dataloader_params=None)

Constructor.

The data module can be used with Path, str or numpy arrays. The data can be either a folder containing images or a single file.

To read custom data types, you can set data_type to custom in data_config and provide a function that returns a numpy array from a path as read_source_func parameter. The function will receive a Path object and an axies string as arguments, the axes being derived from the data_config.

You can also provide a fnmatch and Path.rglob compatible expression (e.g. "*.czi") to filter the files extension using extension_filter.

Parameters:

  • pred_config (InferenceModel) –

    Pydantic model for CAREamics prediction configuration.

  • pred_data (Path or str or ndarray) –

    Prediction data, can be a path to a folder, a file or a numpy array.

  • read_source_func (Callable, default: None ) –

    Function to read custom types, by default None.

  • extension_filter (str, default: '' ) –

    Filter to filter file extensions for custom types, by default "".

  • dataloader_params (dict, default: None ) –

    Dataloader parameters, by default {}.

Raises:

  • ValueError

    If the data type is custom and no read_source_func is provided.

  • ValueError

    If the data type is array and the input is not a numpy array.

  • ValueError

    If the data type is tiff and the input is neither a Path nor a str.

predict_dataloader()

Create a dataloader for prediction.

Returns:

  • DataLoader

    Prediction dataloader.

prepare_data()

Hook used to prepare the data before calling setup.

setup(stage=None)

Hook called at the beginning of predict.

Parameters:

  • stage (Optional[str], default: None ) –

    Stage, by default None.

ProgressBarCallback

Bases: TQDMProgressBar

Progress bar for training and validation steps.

get_metrics(trainer, pl_module)

Override this to customize the metrics displayed in the progress bar.

Parameters:

  • trainer (Trainer) –

    The trainer object.

  • pl_module (LightningModule) –

    The LightningModule object, unused.

Returns:

  • dict

    A dictionary with the metrics to display in the progress bar.

init_test_tqdm()

Override this to customize the tqdm bar for testing.

Returns:

  • tqdm

    A tqdm bar.

init_train_tqdm()

Override this to customize the tqdm bar for training.

Returns:

  • tqdm

    A tqdm bar.

init_validation_tqdm()

Override this to customize the tqdm bar for validation.

Returns:

  • tqdm

    A tqdm bar.

TrainDataModule

Bases: LightningDataModule

CAREamics Ligthning training and validation data module.

The data module can be used with Path, str or numpy arrays. In the case of numpy arrays, it loads and computes all the patches in memory. For Path and str inputs, it calculates the total file size and estimate whether it can fit in memory. If it does not, it iterates through the files. This behaviour can be deactivated by setting use_in_memory to False, in which case it will always use the iterating dataset to train on a Path or str.

The data can be either a folder containing images or a single file.

Validation can be omitted, in which case the validation data is extracted from the training data. The percentage of the training data to use for validation, as well as the minimum number of patches or files to split from the training data can be set using val_percentage and val_minimum_split, respectively.

To read custom data types, you can set data_type to custom in data_config and provide a function that returns a numpy array from a path as read_source_func parameter. The function will receive a Path object and an axies string as arguments, the axes being derived from the data_config.

You can also provide a fnmatch and Path.rglob compatible expression (e.g. "*.czi") to filter the files extension using extension_filter.

Parameters:

  • data_config (DataModel) –

    Pydantic model for CAREamics data configuration.

  • train_data (Path or str or ndarray) –

    Training data, can be a path to a folder, a file or a numpy array.

  • val_data (Path or str or ndarray, default: None ) –

    Validation data, can be a path to a folder, a file or a numpy array, by default None.

  • train_data_target (Path or str or ndarray, default: None ) –

    Training target data, can be a path to a folder, a file or a numpy array, by default None.

  • val_data_target (Path or str or ndarray, default: None ) –

    Validation target data, can be a path to a folder, a file or a numpy array, by default None.

  • read_source_func (Callable, default: None ) –

    Function to read the source data, by default None. Only used for custom data type (see DataModel).

  • extension_filter (str, default: '' ) –

    Filter for file extensions, by default "". Only used for custom data types (see DataModel).

  • val_percentage (float, default: 0.1 ) –

    Percentage of the training data to use for validation, by default 0.1. Only used if val_data is None.

  • val_minimum_split (int, default: 5 ) –

    Minimum number of patches or files to split from the training data for validation, by default 5. Only used if val_data is None.

  • use_in_memory (bool, default: True ) –

    Use in memory dataset if possible, by default True.

Attributes:

  • data_config (DataModel) –

    CAREamics data configuration.

  • data_type (SupportedData) –

    Expected data type, one of "tiff", "array" or "custom".

  • batch_size (int) –

    Batch size.

  • use_in_memory (bool) –

    Whether to use in memory dataset if possible.

  • train_data (Path or ndarray) –

    Training data.

  • val_data (Path or ndarray) –

    Validation data.

  • train_data_target (Path or ndarray) –

    Training target data.

  • val_data_target (Path or ndarray) –

    Validation target data.

  • val_percentage (float) –

    Percentage of the training data to use for validation, if no validation data is provided.

  • val_minimum_split (int) –

    Minimum number of patches or files to split from the training data for validation, if no validation data is provided.

  • read_source_func (Optional[Callable]) –

    Function to read the source data, used if data_type is custom.

  • extension_filter (str) –

    Filter for file extensions, used if data_type is custom.

__init__(data_config, train_data, val_data=None, train_data_target=None, val_data_target=None, read_source_func=None, extension_filter='', val_percentage=0.1, val_minimum_split=5, use_in_memory=True)

Constructor.

Parameters:

  • data_config (DataModel) –

    Pydantic model for CAREamics data configuration.

  • train_data (Path or str or ndarray) –

    Training data, can be a path to a folder, a file or a numpy array.

  • val_data (Path or str or ndarray, default: None ) –

    Validation data, can be a path to a folder, a file or a numpy array, by default None.

  • train_data_target (Path or str or ndarray, default: None ) –

    Training target data, can be a path to a folder, a file or a numpy array, by default None.

  • val_data_target (Path or str or ndarray, default: None ) –

    Validation target data, can be a path to a folder, a file or a numpy array, by default None.

  • read_source_func (Callable, default: None ) –

    Function to read the source data, by default None. Only used for custom data type (see DataModel).

  • extension_filter (str, default: '' ) –

    Filter for file extensions, by default "". Only used for custom data types (see DataModel).

  • val_percentage (float, default: 0.1 ) –

    Percentage of the training data to use for validation, by default 0.1. Only used if val_data is None.

  • val_minimum_split (int, default: 5 ) –

    Minimum number of patches or files to split from the training data for validation, by default 5. Only used if val_data is None.

  • use_in_memory (bool, default: True ) –

    Use in memory dataset if possible, by default True.

Raises:

  • NotImplementedError

    Raised if target data is provided.

  • ValueError

    If the input types are mixed (e.g. Path and numpy.ndarray).

  • ValueError

    If the data type is custom and no read_source_func is provided.

  • ValueError

    If the data type is array and the input is not a numpy array.

  • ValueError

    If the data type is tiff and the input is neither a Path nor a str.

get_data_statistics()

Return training data statistics.

Returns:

  • tuple of list

    Means and standard deviations across channels of the training data.

prepare_data()

Hook used to prepare the data before calling setup.

Here, we only need to examine the data if it was provided as a str or a Path.

TODO: from lightning doc: prepare_data is called from the main process. It is not recommended to assign state here (e.g. self.x = y) since it is called on a single process and if you assign states here then they won't be available for other processes.

https://lightning.ai/docs/pytorch/stable/data/datamodule.html

setup(*args, **kwargs)

Hook called at the beginning of fit, validate, or predict.

Parameters:

  • *args (Any, default: () ) –

    Unused.

  • **kwargs (Any, default: {} ) –

    Unused.

train_dataloader()

Create a dataloader for training.

Returns:

  • Any

    Training dataloader.

val_dataloader()

Create a dataloader for validation.

Returns:

  • Any

    Validation dataloader.

create_careamics_module(algorithm, loss, architecture, use_n2v2=False, struct_n2v_axis='none', struct_n2v_span=5, model_parameters=None, optimizer='Adam', optimizer_parameters=None, lr_scheduler='ReduceLROnPlateau', lr_scheduler_parameters=None)

Create a CAREamics Lightning module.

This function exposes parameters used to create an AlgorithmModel instance, triggering parameters validation.

Parameters:

  • algorithm (SupportedAlgorithm or str) –

    Algorithm to use for training (see SupportedAlgorithm).

  • loss (SupportedLoss or str) –

    Loss function to use for training (see SupportedLoss).

  • architecture (SupportedArchitecture or str) –

    Model architecture to use for training (see SupportedArchitecture).

  • use_n2v2 (bool, default: False ) –

    Whether to use N2V2 or Noise2Void.

  • struct_n2v_axis ("horizontal", "vertical", or "none", default: "none" ) –

    Axis of the StructN2V mask.

  • struct_n2v_span (int, default: 5 ) –

    Span of the StructN2V mask.

  • model_parameters (dict, default: None ) –

    Model parameters to use for training, by default {}. Model parameters are defined in the relevant torch.nn.Module class, or Pyddantic model (see careamics.config.architectures).

  • optimizer (SupportedOptimizer or str, default: 'Adam' ) –

    Optimizer to use for training, by default "Adam" (see SupportedOptimizer).

  • optimizer_parameters (dict, default: None ) –

    Optimizer parameters to use for training, as defined in torch.optim, by default {}.

  • lr_scheduler (SupportedScheduler or str, default: 'ReduceLROnPlateau' ) –

    Learning rate scheduler to use for training, by default "ReduceLROnPlateau" (see SupportedScheduler).

  • lr_scheduler_parameters (dict, default: None ) –

    Learning rate scheduler parameters to use for training, as defined in torch.optim, by default {}.

Returns:

  • CAREamicsModule

    CAREamics Lightning module.

create_predict_datamodule(pred_data, data_type, axes, image_means, image_stds, tile_size=None, tile_overlap=None, batch_size=1, tta_transforms=True, read_source_func=None, extension_filter='', dataloader_params=None)

Create a CAREamics prediction Lightning datamodule.

This function is used to explicitly pass the parameters usually contained in an inference_model configuration.

Since the lightning datamodule has no access to the model, make sure that the parameters passed to the datamodule are consistent with the model's requirements and are coherent. This can be done by creating a Configuration object beforehand and passing its parameters to the different Lightning modules.

The data module can be used with Path, str or numpy arrays. To use array data, set data_type to array and pass a numpy array to train_data.

By default, CAREamics only supports types defined in careamics.config.support.SupportedData. To read custom data types, you can set data_type to custom and provide a function that returns a numpy array from a path. Additionally, pass a fnmatch and Path.rglob compatible expression (e.g. "*.jpeg") to filter the files extension using extension_filter.

In dataloader_params, you can pass any parameter accepted by PyTorch dataloaders, except for batch_size, which is set by the batch_size parameter.

Parameters:

  • pred_data (str or Path or ndarray) –

    Prediction data.

  • data_type ((array, tiff, custom), default: "array" ) –

    Data type, see SupportedData for available options.

  • axes (str) –

    Axes of the data, chosen among SCZYX.

  • image_means (list of float) –

    Mean values for normalization, only used if Normalization is defined.

  • image_stds (list of float) –

    Std values for normalization, only used if Normalization is defined.

  • tile_size (tuple of int, default: None ) –

    Tile size, 2D or 3D tile size.

  • tile_overlap (tuple of int, default: None ) –

    Tile overlap, 2D or 3D tile overlap.

  • batch_size (int, default: 1 ) –

    Batch size.

  • tta_transforms (bool, default: True ) –

    Use test time augmentation, by default True.

  • read_source_func (Callable, default: None ) –

    Function to read the source data, used if data_type is custom, by default None.

  • extension_filter (str, default: '' ) –

    Filter for file extensions, used if data_type is custom, by default "".

  • dataloader_params (dict, default: None ) –

    Pytorch dataloader parameters, by default {}.

Returns:

Notes

If you are using a UNet model and tiling, the tile size must be divisible in every dimension by 2**d, where d is the depth of the model. This avoids artefacts arising from the broken shift invariance induced by the pooling layers of the UNet. If your image has less dimensions, as it may happen in the Z dimension, consider padding your image.

create_train_datamodule(train_data, data_type, patch_size, axes, batch_size, val_data=None, transforms=None, train_target_data=None, val_target_data=None, read_source_func=None, extension_filter='', val_percentage=0.1, val_minimum_patches=5, train_dataloader_params=None, val_dataloader_params=None, use_in_memory=True)

Create a TrainDataModule.

This function is used to explicitly pass the parameters usually contained in a GenericDataConfig to a TrainDataModule.

Since the lightning datamodule has no access to the model, make sure that the parameters passed to the datamodule are consistent with the model's requirements and are coherent.

The default augmentations are XY flip and XY rotation. To use a different set of augmentations, you can pass a list of transforms to transforms.

The data module can be used with Path, str or numpy arrays. In the case of numpy arrays, it loads and computes all the patches in memory. For Path and str inputs, it calculates the total file size and estimate whether it can fit in memory. If it does not, it iterates through the files. This behaviour can be deactivated by setting use_in_memory to False, in which case it will always use the iterating dataset to train on a Path or str.

To use array data, set data_type to array and pass a numpy array to train_data.

By default, CAREamics only supports types defined in careamics.config.support.SupportedData. To read custom data types, you can set data_type to custom and provide a function that returns a numpy array from a path. Additionally, pass a fnmatch and Path.rglob compatible expression (e.g. "*.jpeg") to filter the files extension using extension_filter.

In the absence of validation data, the validation data is extracted from the training data. The percentage of the training data to use for validation, as well as the minimum number of patches to split from the training data for validation can be set using val_percentage and val_minimum_patches, respectively.

In dataloader_params, you can pass any parameter accepted by PyTorch dataloaders, except for batch_size, which is set by the batch_size parameter.

Parameters:

  • train_data (Path or str or ndarray) –

    Training data.

  • data_type ((array, tiff, custom), default: "array" ) –

    Data type, see SupportedData for available options.

  • patch_size (list of int) –

    Patch size, 2D or 3D patch size.

  • axes (str) –

    Axes of the data, chosen amongst SCZYX.

  • batch_size (int) –

    Batch size.

  • val_data (Path or str or ndarray, default: None ) –

    Validation data, by default None.

  • transforms (list of Transforms, default: None ) –

    List of transforms to apply to training patches. If None, default transforms are applied.

  • train_target_data (Path or str or ndarray, default: None ) –

    Training target data, by default None.

  • val_target_data (Path or str or ndarray, default: None ) –

    Validation target data, by default None.

  • read_source_func (Callable, default: None ) –

    Function to read the source data, used if data_type is custom, by default None.

  • extension_filter (str, default: '' ) –

    Filter for file extensions, used if data_type is custom, by default "".

  • val_percentage (float, default: 0.1 ) –

    Percentage of the training data to use for validation if no validation data is given, by default 0.1.

  • val_minimum_patches (int, default: 5 ) –

    Minimum number of patches to split from the training data for validation if no validation data is given, by default 5.

  • train_dataloader_params (dict, default: None ) –

    Pytorch dataloader parameters for the training data, by default {}.

  • val_dataloader_params (dict, default: None ) –

    Pytorch dataloader parameters for the validation data, by default {}.

  • use_in_memory (bool, default: True ) –

    Use in memory dataset if possible, by default True.

Returns:

Examples:

Create a TrainingDataModule with default transforms with a numpy array:

>>> import numpy as np
>>> from careamics.compat.lightning import create_train_datamodule
>>> my_array = np.arange(256).reshape(16, 16)
>>> data_module = create_train_datamodule(
...     train_data=my_array,
...     data_type="array",
...     patch_size=(8, 8),
...     axes='YX',
...     batch_size=2,
... )

For custom data types (those not supported by CAREamics), then one can pass a read function and a filter for the files extension:

>>> import numpy as np
>>> from careamics.compat.lightning import create_train_datamodule
>>>
>>> def read_npy(path):
...     return np.load(path)
>>>
>>> data_module = create_train_datamodule(
...     train_data="path/to/data",
...     data_type="custom",
...     patch_size=(8, 8),
...     axes='YX',
...     batch_size=2,
...     read_source_func=read_npy,
...     extension_filter="*.npy",
... )

If you want to use a different set of augmentations, you can pass a list of transforms:

>>> import numpy as np
>>> from careamics.compat.lightning import create_train_datamodule
>>> from careamics.config.augmentations import XYFlipConfig
>>> from careamics.config.support import SupportedTransform
>>> my_array = np.arange(256).reshape(16, 16)
>>> my_transforms = [
...     XYFlipConfig(flip_y=False),
... ]
>>> data_module = create_train_datamodule(
...     train_data=my_array,
...     data_type="array",
...     patch_size=(8, 8),
...     axes='YX',
...     batch_size=2,
...     transforms=my_transforms,
... )