dgs.utils.torchtools.save_checkpoint

dgs.utils.torchtools.save_checkpoint(state: dict[str, any], save_dir: str, *, is_best: bool = False, remove_module_from_keys: bool = False, prepend: str = '', verbose: bool = True) → None[source]

Save a given checkpoint.

Parameters:

  • state – State dictionary. See examples.

  • save_dir – directory to save checkpoint.

  • is_best (bool, optional) – if True, this checkpoint will be copied and named model-best.pth.tar. Default is False.

  • remove_module_from_keys – Whether to remove the ‘module.’ prepend in the state dict of the module.

  • prepend – A string to prepend to the filename.

  • verbose (bool, optional) – whether to print a confirmation when the checkpoint has been created. Default is True.

Examples

>>> state = {
>>>     'model': model.state_dict(),
>>>     'epoch': 10,
>>>     'rank1': 0.5,
>>>     'optimizer': optimizer.state_dict()
>>> }
>>> save_checkpoint(state, 'log/my_model')