IO

load_checkpoint

load_checkpoint(file, weights_only=False, **kwargs)[source]

Loads a saved model

Parameters:
  • file (str) – filepath to a file containing a saved model
  • weights_only (bool) – whether the file contains only weights / only weights should be returned
  • **kwargs – Additional keyword arguments (passed to torch.load) Especially “map_location” is important to change the device the state_dict should be loaded to
Returns:

  • OrderedDict – checkpoint state_dict if weights_only=True
  • torch.nn.Module, OrderedDict, int – Model, Optimizers, epoch with loaded state_dicts if weights_only=False

save_checkpoint

save_checkpoint(file: str, model=None, optimizers={}, epoch=None, weights_only=False, **kwargs)[source]

Save model’s parameters

Parameters:
  • file (str) – filepath the model should be saved to
  • model (AbstractNetwork or None) – the model which should be saved if None: empty dict will be saved as state dict
  • optimizers (dict) – dictionary containing all optimizers
  • epoch (int) – current epoch (will also be pickled)
  • weights_only (bool) – whether or not to save only the model’s weights or also save additional information (for easy loading)