manugloeck/dmri-pretrained

DMRI pretrained checkpoints

Pretrained checkpoints for mackelab/dmri, stored as portable Orbax bundles with their frozen Hydra configuration.

Model	Width	Best step	Bundle size
b3s_2_4_6_64	64	2,640,000	55 MB
b3s_2_4_6_128	128	1,800,000	123 MB

Load a model

from flax import nnx
from dmri.train.utils import load_checkpoint

checkpoint, model, simulators = load_checkpoint(
    repo_id="manugloeck/dmri-pretrained",
    model_name="b3s_2_4_6_64",
    which="best",
)
params = checkpoint.get("params_ema", checkpoint["params"])
nnx.update(model, params)
model.eval()

Only the requested model subfolder is downloaded and cached by huggingface_hub.