Skip to content
  • API documentation
  • careamics
  • lightning
  • callbacks
  • careamics_checkpoint_info_callback

careamics_checkpoint_info_callback

Lightning callback for storing CAREamics configuration in checkpoints.

CareamicsCheckpointInfo #

Bases: Callback

Callback to save CAREamics configuration in Lightning checkpoints.

This callback automatically stores CAREamics version, experiment name, and training configuration in the checkpoint file for reproducibility.

Parameters:

Name Type Description Default
careamics_version str

Version of CAREamics used for training.

 required
experiment_name str

Name of the experiment.

 required
training_config TrainingConfig

Training configuration to store in checkpoint.

 required

Attributes:

Name Type Description
careamics_version str

Version of CAREamics used for training.
experiment_name str

Name of the experiment.
training_config TrainingConfig

Training configuration to store in checkpoint.
Source code in src/careamics/lightning/callbacks/careamics_checkpoint_info_callback.py 
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
class CareamicsCheckpointInfo(Callback):
    """
    Callback to save CAREamics configuration in Lightning checkpoints.

    This callback automatically stores CAREamics version, experiment name,
    and training configuration in the checkpoint file for reproducibility.

    Parameters
    ----------
    careamics_version : str
        Version of CAREamics used for training.
    experiment_name : str
        Name of the experiment.
    training_config : TrainingConfig
        Training configuration to store in checkpoint.

    Attributes
    ----------
    careamics_version : str
        Version of CAREamics used for training.
    experiment_name : str
        Name of the experiment.
    training_config : TrainingConfig
        Training configuration to store in checkpoint.
    """

    def __init__(
        self,
        careamics_version: str,
        experiment_name: str,
        training_config: TrainingConfig,
    ):
        """
        Initialize the callback.

        Parameters
        ----------
        careamics_version : str
            Version of CAREamics used for training.
        experiment_name : str
            Name of the experiment.
        training_config : TrainingConfig
            Training configuration to store in checkpoint.
        """
        super().__init__()
        self.careamics_version = careamics_version
        self.experiment_name = experiment_name
        self.training_config = training_config

    def on_save_checkpoint(
        self, trainer: Trainer, pl_module: LightningModule, checkpoint: dict[str, Any]
    ) -> None:
        """
        Lightning hook called when saving a checkpoint.

        Adds CAREamics configuration to the checkpoint dictionary.

        Parameters
        ----------
        trainer : Trainer
            Lightning trainer instance.
        pl_module : LightningModule
            Lightning module being trained.
        checkpoint : dict[str, Any]
            Checkpoint dictionary to modify.
        """
        checkpoint["careamics_info"] = {
            "version": self.careamics_version,
            "experiment_name": self.experiment_name,
            "training_config": self.training_config.model_dump(mode="json"),
        }

__init__(careamics_version, experiment_name, training_config) #

Initialize the callback.

Parameters:

Name Type Description Default
careamics_version str

Version of CAREamics used for training.

 required
experiment_name str

Name of the experiment.

 required
training_config TrainingConfig

Training configuration to store in checkpoint.

 required
Source code in src/careamics/lightning/callbacks/careamics_checkpoint_info_callback.py 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
def __init__(
    self,
    careamics_version: str,
    experiment_name: str,
    training_config: TrainingConfig,
):
    """
    Initialize the callback.

    Parameters
    ----------
    careamics_version : str
        Version of CAREamics used for training.
    experiment_name : str
        Name of the experiment.
    training_config : TrainingConfig
        Training configuration to store in checkpoint.
    """
    super().__init__()
    self.careamics_version = careamics_version
    self.experiment_name = experiment_name
    self.training_config = training_config

on_save_checkpoint(trainer, pl_module, checkpoint) #

Lightning hook called when saving a checkpoint.

Adds CAREamics configuration to the checkpoint dictionary.

Parameters:

Name Type Description Default
trainer Trainer

Lightning trainer instance.

 required
pl_module LightningModule

Lightning module being trained.

 required
checkpoint dict[str, Any]

Checkpoint dictionary to modify.

 required
Source code in src/careamics/lightning/callbacks/careamics_checkpoint_info_callback.py 
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
def on_save_checkpoint(
    self, trainer: Trainer, pl_module: LightningModule, checkpoint: dict[str, Any]
) -> None:
    """
    Lightning hook called when saving a checkpoint.

    Adds CAREamics configuration to the checkpoint dictionary.

    Parameters
    ----------
    trainer : Trainer
        Lightning trainer instance.
    pl_module : LightningModule
        Lightning module being trained.
    checkpoint : dict[str, Any]
        Checkpoint dictionary to modify.
    """
    checkpoint["careamics_info"] = {
        "version": self.careamics_version,
        "experiment_name": self.experiment_name,
        "training_config": self.training_config.model_dump(mode="json"),
    }