pytorch_pfn_extras.training.extensions.snapshot_object#

pytorch_pfn_extras.training.extensions.snapshot_object(target, filename, savefun=None, *, condition=None, writer=None, snapshot_on_error=False, n_retains=-1, autoload=False)#

Returns an extension to take snapshots of a given object.

This extension serializes the given object and saves it to the output directory.

This extension is called once per epoch by default. To take a snapshot at a different interval, a trigger object specifying the required interval can be passed along with this extension to the extend() method of the manager.

The default priority is lower than that of most built-in extensions.

Parameters:

  • target (Any) – Object to serialize.

  • filename (str) – Name of the file into which the object is serialized. It can be a format string, where the manager object is passed to the str.format() method. For example, 'snapshot_{.iteration}' is converted to 'snapshot_10000' at the 10,000th iteration.

  • savefun (Optional[Any]) – Function to save the object. It takes two arguments: the output file path and the object to serialize.

  • condition – Condition object. It must be a callable object that returns boolean without any arguments. If it returns True, the snapshot will be done. If not, it will be skipped. The default is a function that always returns True.

  • writer – Writer object. It must be a callable object. See below for the list of built-in writers. If savefun is other than None, this argument must be None. In that case, a SimpleWriter object instantiated with specified savefun argument will be used.

  • snapshot_on_error (bool) – Whether to take a snapshot in case training loop has failed.

  • n_retains (int) – Number of snapshot files to retain through the cleanup. Must be a positive integer for any cleanup to take place. Automatic deletion of old snapshots only works when the filename is string.

  • autoload (bool) – With this enabled, the extension automatically finds the latest snapshot and loads the data to the target. Automatic loading only works when the filename is a string.

  • saver_rank (int) – If defined, the snapshot will be taken by only one rank when running in distributed mode and restored by all.

  • kwargs (Any) –

Returns:

Snapshot extension object.

Return type:

_Snapshot

See also