Training Python API¶

The supported training interface is the dmri train command. The functions below are reusable when constructing simulators, models, dataloaders, or checkpoint bundles in Python. The functions and state classes in dmri.train.train_script implement the CLI loop and are not a supported programmatic trainer or callback interface.

See the training guide for configuration, output, resume, and evaluation examples.

Builders¶

build_simulator ¶

build_simulator(cfg: DictConfig)

Build the signal-model class and training-data generators.

build_model ¶

build_model(cfg: DictConfig, sim_type)

build_simulator(cfg) returns the configured MultiCompartment class and one training-data generator per acquisition entry. build_model(cfg, sim_type) builds an initialized model for that simulator class.

Simulation data¶

SimulationDataset ¶

SimulationDataset(simulator_fn: Callable[..., SimOutput], *, simulation_batch_size: int = 128, rng: RngKey, simulation_device: Device, jit_simulator: bool = True, return_numpy: bool = False, buffer_size: int = 8192)

Fixed-size dataset whose samples are refreshed asynchronously.

reset ¶

reset(*, seed: int | None = None, rng: RngKey | None = None) -> None

close ¶

close() -> None

Stop the producer thread and clean up resources.

set_data ¶

set_data(data: Any) -> None

Replace the internal buffer with user-provided data.

Parameters¶

data : Any A PyTree of arrays with a leading sample dimension. Leaves must be array-like and broadcast-consistent in their first dimension. start_producer : bool, default False If True, (re)start the background producer after setting the buffer. By default we keep the dataset static and the producer stopped.

get_stats ¶

get_stats() -> dict[str, Any]

instantiate_dataloader ¶

instantiate_dataloader(dataset: Any, loader_cfg: Any, *, seed: int | None = None, default_shuffle: bool | None = None, default_drop_last: bool | None = None) -> DataLoader

Instantiate a probjax DataLoader with normalised parameters.

Checkpoint loading¶

load_cfg ¶

load_cfg(path)

load_checkpoint ¶

load_checkpoint(path=None, which='latest', partial_restore=False, *, repo_id=None, model_name=None, revision=None, cache_dir=None, token=None, local_files_only=False, precision=None)

Restore a checkpoint from a local result directory or Hugging Face.

Provide path for a local checkpoint. For a remote bundle, provide both repo_id and model_name; the selected subfolder is downloaded into the Hugging Face cache before the normal Orbax restore path is used.

precision selects the compute dtype for the rebuilt model (see :mod:dmri.eval.precision); the default keeps the checkpoint's own setting.

load_checkpoint accepts a local run or a Hugging Face model subfolder. It returns (checkpoint, model, simulator_class). The model is rebuilt in eval mode, but the returned checkpoint parameters are not applied automatically.

Checkpoint bundles¶

bundle_checkpoint ¶

bundle_checkpoint(path, output_dir, model_name=None, which='best_and_latest') -> Path

Create a portable checkpoint bundle for a model repository.

The bundle contains a root config.yaml and an Orbax checkpoint tree, so it can be restored without the original timestamped Hydra run directory.

Parameters:

Name	Description	Default
`path`	Local training-result directory.	required
`output_dir`	Directory in which to create the model subfolder.	required
`model_name`	Subfolder name. Defaults to `cfg.name`.	`None`
`which`	`"best_and_latest"` (default), `"best"`, `"latest"`, `"all"`, or an integer step.	`'best_and_latest'`

Returns:

Type	Description
`Path`	Path to the created model subfolder.

upload_checkpoint_to_hub ¶

upload_checkpoint_to_hub(path, repo_id, *, model_name=None, which='best_and_latest', private=False, token=None, commit_message=None)

Upload a local run into a model subfolder in one Hugging Face repo.

Authentication uses the cached Hugging Face token by default. Pass token explicitly for non-interactive or private-repository workflows.

Returns:

Type	Description
	The Hugging Face commit information returned by `upload_folder`.

download_checkpoint_from_hub ¶

download_checkpoint_from_hub(repo_id, model_name, *, revision=None, cache_dir=None, token=None, local_files_only=False) -> Path

Download one model subfolder from a Hugging Face checkpoint repo.

Only <model_name>/** is downloaded, allowing several variants to share one repository without downloading every checkpoint.