Reference / API reference / dorothy / dorothy.inference / dorothy.inference.checkpoint

dorothy.inference.checkpoint

Load trained DOROTHY checkpoints.

A checkpoint directory written by Trainer.save_checkpoint() contains the model weights, config.yaml, (usually) normalizer.pkl and the calibration factors in calibration.json, if training fit them. The model is rebuilt from the saved config, so every architecture option matches training. Directories with weights only fall back to inferring a plain MLP, with an AdditiveOutput layer, from the state dict.

Functions

build_model(config)	Build the untrained model described by an experiment config.
head_label_sources(model, config)	Label source of each output head, in model order.
infer_mlp_architecture(state_dict)	Infer MLP constructor arguments from a state dict.
load_checkpoint(checkpoint_path[, ...])	Load a trained model, its normalizer, config and calibration factors.

Classes

LoadedCheckpoint(model, normalizer, config, ...)

A trained model with the artifacts needed to use it.

class dorothy.inference.checkpoint.LoadedCheckpoint(model, normalizer, config, path, device, calibration)

Bases: object

A trained model with the artifacts needed to use it.

model

The model, in eval mode on the requested device.

normalizer

Label normalizer, or None if the checkpoint has none.

config

The experiment config the model was trained with, or None for weights-only checkpoints.

path

The checkpoint directory.

device

Device the model was moved to.

calibration

Calibration factors of best_model.pth, label source -> parameter name -> factor. None without calibration.json, without config.yaml, or for another weights file.

property parameter_names: list[str]

Names of the predicted parameters, in output order.

property scatter_floor: float

Scatter floor the model was trained with (0.01 without a config).

dorothy.inference.checkpoint.infer_mlp_architecture(state_dict)

Infer MLP constructor arguments from a state dict.

Recovers input/output sizes, hidden layer sizes (in layer order) and the normalization type. The activation cannot be inferred from weights and is left at the MLP default.

Parameters:

state_dict (dict[str, Tensor]) – MLP state dict.

Returns:

Dict of keyword arguments for MLP.

Raises:

ValueError – If the state dict has no linear layers, or has a label-aware output layer (its config is needed to rebuild it).

Return type:

dict

dorothy.inference.checkpoint.build_model(config)

Build the untrained model described by an experiment config.

Parameters:

config (ExperimentConfig) – The experiment config.

Returns:

A MultiHeadMLP if the config is multi-head, else an MLP.

Return type:

MLP | MultiHeadMLP

dorothy.inference.checkpoint.head_label_sources(model, config)

Label source of each output head, in model order.

A plain MLP’s head is named after the configured label source.

Parameters:

• model (MLP | MultiHeadMLP) – The model.

• config (ExperimentConfig) – The experiment config the model was built from.

Returns:

One label source name per output head.

Return type:

list[str]

dorothy.inference.checkpoint.load_checkpoint(checkpoint_path, model_file='best_model.pth', normalizer_file='normalizer.pkl', device='auto')

Load a trained model, its normalizer, config and calibration factors.

Parameters:

• checkpoint_path (str | Path) – Checkpoint directory.

• model_file (str) – Weights file inside the directory. The calibration factors in calibration.json describe best_model.pth, so other files load without them.

• normalizer_file (str | None) – Normalizer file name (None to skip loading it).

• device (str | device) – Device for the model (“auto”, “cpu”, “cuda”, “mps”).

Returns:

LoadedCheckpoint with the model in eval mode.

Raises:

• FileNotFoundError – If the weights file does not exist.

• ValueError – If a checkpoint without config.yaml holds a model with label-aware uncertainty.

Return type:

LoadedCheckpoint