CkptMixin

class torch_ecg.utils.CkptMixin[source]

Bases: object

Mixin class for loading from checkpoint class methods

classmethod from_checkpoint(path: Union[str, pathlib.Path], device: Optional[torch.device] = None) → Tuple[torch.nn.modules.module.Module, dict][source]

Load a model from a checkpoint.

Parameters

  • path (str or pathlib.Path) – 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: Union[str, pathlib.Path], train_config: torch_ecg.cfg.CFG) → None[source]

Save the model to disk.

Parameters

  • path (str or pathlib.Path) – Path to save the model.

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

Return type

None