CheckpointIO

class lightning.pytorch.plugins.io.CheckpointIO

Bases: ABC

Interface to save/load checkpoints as they are saved through the Strategy.

Warning

This is an experimental feature.

Typically most plugins either use the Torch based IO Plugin; TorchCheckpointIO but may require particular handling depending on the plugin.

In addition, you can pass a custom CheckpointIO by extending this class and passing it to the Trainer, i.e Trainer(plugins=[MyCustomCheckpointIO()]).

Note

For some plugins, it is not possible to use a custom checkpoint plugin as checkpointing logic is not modifiable.

abstract load_checkpoint(path, *, state=None, map_location=None, weights_only=None)

Load checkpoint from a path when resuming or loading ckpt for test/validate/predict stages.

Parameters:

• path (Union[str, Path]) – Path to checkpoint

• state (Union[Module, Optimizer, dict[str, Union[Module, Optimizer, Any]], None]) – Optional dict to load the checkpoint into.

• map_location (Optional[Any]) – a function, torch.device, string or a dict specifying how to remap storage locations.

• weights_only (Optional[bool]) – Defaults to None. If True, restricts loading to state_dicts of plain torch.Tensor and other primitive types. If loading a checkpoint from a trusted source that contains an nn.Module, use weights_only=False. If loading checkpoint from an untrusted source, we recommend using weights_only=True. For more information, please refer to the PyTorch Developer Notes on Serialization Semantics.

Return type:

dict[str, Any]

Returns:

• A dictionary containing checkpoint contents that still need to be

  applied by the caller, or

• None if the checkpoint was fully restored in-place into state.

abstract remove_checkpoint(path)

Remove checkpoint file from the filesystem.

Parameters:

path (Union[str, Path]) – Path to checkpoint

Return type:

None

abstract save_checkpoint(checkpoint, path, storage_options=None)

Save model/training states as a checkpoint file through state-dump and file-write.

Parameters:

• checkpoint (dict[str, Any]) – dict containing model and trainer state

• path (Union[str, Path]) – write-target path

• storage_options (Optional[Any]) – Optional parameters when saving the model/training states.

Return type:

None

teardown()

This method is called to teardown the process.

Return type:

None

property requires_state_on_load: bool

Whether the state argument of load_checkpoint is required for loading the checkpoint.

If True, the Trainer will always pass a state dict containing the current model and optimizer to the load_checkpoint method. This is for plugins that need to do in-place restoration of the checkpoint into the provided state objects instead of returning a new checkpoint dict.