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) → Tuple[Module, dict][source]

Load a model from a checkpoint.

Parameters:

  • path (path-like) – Path of the checkpoint.

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

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[source]

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