Ctrl+K

nvalchemi.training.load_checkpoint#

nvalchemi.training.load_checkpoint(root_folder, checkpoint_index=-1, map_location=None, model_names=None, *, adapter=None, adapter_kwargs=None, validators=None, hooks=None, training_fn=None, strategy=None)[source]#

Load a multi-component checkpoint written by save_checkpoint().

Components are rebuilt in dependency order: models first, then optimizers (which need model parameters), then schedulers (which need an optimizer instance). Associations from the manifest wire each optimizer to the correct model and each scheduler to the correct optimizer.

Parameters:

  • root_folder (Path | str) – Root directory containing manifest.json.

  • checkpoint_index (int) – Index of the checkpoint to load. -1 (default) loads the latest index recorded in the manifest.

  • map_location (str | device | None) – Forwarded to every torch.load() call. When not None, each loaded model is additionally moved via model.to(map_location). Optimizers and schedulers have their state placed by torch.load alone (they lack a standard .to() API).

  • model_names (Iterable[str] | None) – If given, load only the models with these names together with the optimizers and schedulers wired to them through manifest.associations. Accepts any iterable of strings (typically a set). None (default) loads every component on disk. The returned manifest’s associations still reflects the full on-disk mapping, so callers can inspect what was not loaded.

  • adapter (str | None) – Optional foreign-checkpoint adapter name. V1 supports "mace" for trusted local MACE .pt files.

  • adapter_kwargs (Mapping[str, Any] | None) – Adapter-specific options. For adapter="mace", accepted keys are model_name, dtype, enable_cueq, compile_model, and compile_kwargs.

  • validators (Sequence[Callable[[str, Mapping[str, Any], Mapping[str, Any]], None]] | None) – Optional callbacks invoked as validator(model_name, entry, loaded) for each high-level loaded model entry. Use these for model-specific chemistry or topology compatibility checks.

  • hooks (Sequence[Any] | None) – Runtime hooks supplied when reconstructing a saved strategy.

  • training_fn (Any) – Runtime training function override supplied when reconstructing a saved strategy.

  • strategy (Any | None) – Optional already-constructed strategy to hydrate from the checkpoint. This mode restores model, optimizer, scheduler, runtime-counter, and checkpointable hook state into the live objects instead of rebuilding models from saved specs.

Returns:

  • CheckpointManifest – For legacy component-only checkpoints, a hydrated manifest is returned.

  • dict[str, Any] – For strategy checkpoints or adapter loads, a builtin dict containing strategy, models, manifest, checkpoint_index, and source is returned.

Raises:

  • FileNotFoundError – If manifest.json is missing or a checkpoint .pt file does not exist.

  • KeyError – If any name in model_names does not appear in manifest.models.

  • RuntimeError – If a model spec does not build an Module.

Return type:

CheckpointManifest | dict[str, Any]

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)})
...     result = load_checkpoint(tmp)
...     isinstance(result.models["main"][0], nn.Linear)
True

Loading onto CPU regardless of the original device:

result = load_checkpoint("runs/exp1", map_location="cpu")

Selecting a subset of models (e.g., teacher and student but not the third auxiliary model):

result = load_checkpoint("runs/kd", model_names={"teacher", "student"})
On this page
Show Source