TorchCheckpointIO

class lightning.pytorch.plugins.io.TorchCheckpointIO

기반 클래스: lightning.fabric.plugins.io.checkpoint_io.CheckpointIO

CheckpointIO that utilizes torch.save() and torch.load() to save and load checkpoints respectively, common for most use cases.

경고

This is an experimental feature.

load_checkpoint(path, map_location=<function TorchCheckpointIO.<lambda>>)

Loads checkpoint using torch.load(), with additional handling for fsspec remote loading of files.

매개변수

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

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

Returns: The loaded checkpoint.

예외 발생

FileNotFoundError – If path is not found by the fsspec filesystem

반환 형식

Dict[str, Any]

remove_checkpoint(path)

Remove checkpoint file from the filesystem.

매개변수

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

반환 형식

None

save_checkpoint(checkpoint, path, storage_options=None)

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

매개변수

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

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

• storage_options (Optional[Any]) – not used in TorchCheckpointIO.save_checkpoint

예외 발생

TypeError – If storage_options arg is passed in

반환 형식

None