xnn.common.models.hub.checkpoint#

Reading, writing and building models from checkpoints.

Three kinds of file hold a model’s weights, and load_checkpoint() reads all of them into the same Checkpoint:

  • a portable model directory (card.json, config.yaml, model.pt), written by save_pretrained(). The config is plain YAML and the weights load with torch.load(weights_only=True), so the directory carries no pickled Python objects and no absolute paths: copy it anywhere and it loads there;

  • an xnn trainer checkpoint (best.pt: {"model": state_dict, "cfg": Config}, written by save()), or a bare state dict;

  • a file in a foreign format registered in formats (e.g. a mace-torch foundation model), converted on the fly.

build_potential() turns a checkpoint into a ready model. It is the one loading recipe behind from_pretrained(), the MDI engine, xnn export and the benchmark runner.

Functions

build_potential(config, state_dict, *[, ...])

Rebuild a model from its config and weights.

config_to_dict(config)

Serialize a Config (or a bare ModelConfig) to a dict.

core_model(model)

The innermost potential under the force, dispersion, LES and multi-head wrappers.

file_references(config)

Config values of the model section that name existing files.

float_dtype(source)

Floating-point dtype of a module's parameters or of a state dict.

has_dispersion(model)

Whether a model already includes a D3 / D4 wrapper (at any nesting).

load_checkpoint(path, *[, format, head])

Read any supported checkpoint into a Checkpoint.

read_config(path)

Read a config written by write_config().

save_pretrained(model, save_directory, *[, ...])

Write a model as a portable directory.

wrapped_state_dict(model)

State dict in the trainer's ForceStressOutput layout, on CPU.

write_config(config, path)

Write a config as YAML, atomically (see config_to_dict()).

Classes

Checkpoint(config, state_dict, card, path)

A model's weights and the config that rebuilds it.
Show Source