CkptMixin#

class torch_ecg.utils.CkptMixin[source]#

Bases: object

Mixin class for loading from checkpoint class methods

classmethod from_checkpoint(path: str | bytes | PathLike, device: device | None = None, weights_only: Literal[True, False, 'auto'] = 'auto') → Tuple[Module, dict][source]#

Load a model from a checkpoint.

Parameters:

  • path (path-like) – Path to the checkpoint. If it is a directory, then this directory should be one of the following cases: - contain a model.safetensors file, a model_config.json file, and a train_config.json file. - contain only one checkpoint file (with the extension .pth, .pt, .safetensors).

  • device (torch.device, optional) – Map location of the model parameters, defaults to “cuda” if available, otherwise “cpu”.

  • weights_only ({"auto", True, False}, default "auto") –

    Whether to load only the weights of the model.

    Added in version 0.0.31.

Returns:

  • model (torch.nn.Module) – The model loaded from a checkpoint.

  • aux_config (dict) – Auxiliary configs that are needed for data preprocessing, etc.

classmethod from_remote(url: str, model_dir: str | bytes | PathLike, filename: str | None = None, device: device | None = None, weights_only: Literal[True, False, 'auto'] = 'auto') → Tuple[Module, dict][source]#

Load the model from the remote model.

Parameters:

  • url (str) – URL of the remote model.

  • model_dir (path-like) – Path for downloading the model.

  • filename (str, optional) – Filename of the model to save, defaults to the basename of the URL.

  • device (torch.device, optional) – Map location of the model parameters, defaults to “cuda” if available, otherwise “cpu”.

  • weights_only ({"auto", True, False}, default "auto") –

    Whether to load only the weights of the model.

    Added in version 0.0.31.

Returns:

  • model (torch.nn.Module) – The model loaded from a checkpoint.

  • aux_config (dict) – Auxiliary configs that are needed for data preprocessing, etc.

save(path: str | bytes | PathLike, train_config: CFG, extra_items: dict | None = None, use_safetensors: bool | None = None, safetensors_single_file: bool = True) → None[source]#

Save the model to disk.

Note

If path has a suffix of .pth or .pt, safetensors will NOT be used unless use_safetensors=True is explicitly set. Otherwise, safetensors is used by default.

Parameters:

  • path (path-like) – Path to save the model.

  • train_config (CFG) – Config for training the model, used when one restores the model.

  • extra_items (dict, optional) –

    Extra items to save along with the model. The values should be serializable: can be saved as a json file, or is a dict of torch tensors.

    Added in version 0.0.32.

  • use_safetensors (bool, optional) –

    Whether to use safetensors to save the model. If None, it is determined by the suffix of path: if it is .pth or .pt, then use_safetensors is set to False; otherwise it is set to True. If True, and the suffix of path is .pth or .pt, the suffix is changed to .safetensors.

    Added in version 0.0.32.

  • safetensors_single_file (bool, default True) –

    Whether to save the metadata along with the state dict into one file.

    Added in version 0.0.32.

Return type:

None