CkptMixin

class torch_ecg.utils.CkptMixin

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]

Load a model from a checkpoint.

Parameters:

• path (path-like) – Path to the checkpoint. If it is a directory, then this directory should contain only one checkpoint file (with the extension .pth or .pt).

• 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]

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) → None

Save the model to disk.

Parameters:

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

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

Return type:

None