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 toTrue
.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 toTrue
, thus returning the unaltered model.
See also
- __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