nvalchemi.training.save_checkpoint#

nvalchemi.training.save_checkpoint(root_folder, models=None, optimizers=None, schedulers=None, associations=None, checkpoint_index=-1, strategy=None)[source]#

Save a checkpoint with a manifest.

The low-level component form accepts explicit models, optimizers, and schedulers mappings. The strategy-aware form accepts strategy=TrainingStrategy(...) (or the strategy as the second positional argument) and writes additional strategy.json metadata with the serializable recipe and restart counters.

Parameters:

root_folder (Path | str) – Root directory for the checkpoint tree.
models (dict[str, tuple[Module, BaseSpec]] | Any | None) – Mapping of model name to (module, spec) pairs, or a TrainingStrategy instance.
optimizers (dict[str, tuple[Optimizer, BaseSpec]] | None) – Optional mapping of optimizer name to (optimizer, spec) pairs.
schedulers (dict[str, tuple[LRScheduler, BaseSpec]] | None) – Optional mapping of scheduler name to (scheduler, spec) pairs.
associations (dict[str, dict[str, Any]] | None) – Optional model-centric linkage mapping a model name to {"optimizers": [...], "schedulers": [...]}. When None (default), associations are inferred automatically by matching optimizer param_groups to model parameters via data_ptr() identity, and schedulers to optimizers via object identity.
checkpoint_index (int) – Index for the checkpoint files. -1 (default) auto-increments from the manifest’s last index, or starts at 0.
strategy (Any | None) – Optional training strategy to save as a restartable checkpoint.

Returns:

The checkpoint index that was written.

Return type:

int

Raises:

ValueError – If an existing spec.json disagrees with the spec being saved (ignoring timestamp).

Examples

>>> import tempfile, torch.nn as nn
>>> from nvalchemi.training._spec import create_model_spec
>>> with tempfile.TemporaryDirectory() as tmp:
...     spec = create_model_spec(nn.Linear, in_features=4, out_features=2)
...     save_checkpoint(tmp, models={"main": (nn.Linear(4, 2), spec)})
0