io

Save and load (the state of) PyTorch models to and from disk.

class ModelLoader(path='', map_location=None)[source]

Bases: ArgRepr

Load a previously saved model from file.

Parameters:
  • path (str, optional) – Full or partial path to the model to load. If not fully specified here, it can be completed on calling the instance. Defaults to the current working directory of the python interpreter.

  • map_location (str or device, optional) – The device to load the modelo onto. Defaults to None which loads to the PyTorch default device.

__call__(path='')[source]

Load a previously saved model from file.

Parameters:

path (str, optional) – Path (including file name) to the model file to load. If it starts with a backslash, it will be interpreted as absolute, if not, as relative to the path specified at instantiation. Defaults to an empty string, which results in an unchanged path.

Returns:

The loaded model.

Return type:

Module

class ModelSaver(path='', create=False)[source]

Bases: ArgRepr

Save an entire model to file.

Parameters:
  • path (str, optional) – Path (including file name) to save the model to. May include any number of string placeholders (i.e., pairs of curly brackets) that will be interpolated when instances are called. Defaults to the current working directory of the python interpreter.

  • create (bool, optional) – What to do if the directory where the model should be saved does not exist. Defaults to False.

__call__(model, *parts)[source]

Save a model to file.

Parameters:
  • model (Module) – The model to save.

  • *parts (str, optional) – Fragments that will be interpolated into the path string given at instantiation. Obviously, there must be at least as many as there are placeholders in the path.

Returns:

An empty tuple.

Return type:

tuple

class StateLoader(path='', map_location=None, merge=True, not_found=NotFound.RAISE)[source]

Bases: ArgRepr

Load the state of a model from file.

Parameters:
  • path (str, optional) – Path (including file name) to the model’s state_dict() on disk. May include any number of string placeholders (i.e., pairs of curly brackets) that will be interpolated when instances are called. Defaults to the current working directory of the python interpreter.

  • map_location (str or Device, optional) – The device to load the state onto. Defaults to None which loads to the PyTorch device(s) that were saved with the model.

  • merge (bool, optional) – Whether the loaded state should be merged into the state of the model (True) or replace its state (False). Defaults to True.

  • not_found (str, optional) – What to do if the specified file is not found. One of “ignore”, “warn”, or “raise” (use the NotFound enum to avoid typos). Defaults to “raise”. If set to “ignore” or “warn” and the specified file is not found, merge is overridden to True, thus returning the unaltered model.

See also

NotFound

__call__(model, *parts)[source]

Load the state of a model from file.

Parameters:
  • model (Module) – Model to load the state of.

  • *parts (str, optional) – Fragments that will be interpolated into the path string given at instantiation. Obviously, there must be at least as many as there are placeholders in the path.

Returns:

The model with its state restored.

Return type:

Module

class StateSaver(path='', create=False)[source]

Bases: ArgRepr

Save the state of a model to file.

Parameters:
  • path (str, optional) – Path (including file name) to save a model’s state_dict() to. May include any number of string placeholders (i.e., pairs of curly brackets) that will be interpolated when instances are called. Defaults to the current working directory of the python interpreter.

  • create (bool, optional) – What to do if the directory where the state should be saved does not exist. Defaults to False.

__call__(model, *parts)[source]

Save the state of a model, optimizer, or scheduler to file.

Parameters:
  • model (Module) – Model to save the state of.

  • *parts (str, optional) – Fragments that will be interpolated into the path string given at instantiation. Obviously, there must be at least as many as there are placeholders in the path.

Returns:

An empty tuple.

Return type:

tuple