Utils

Inference contracts, Safetensors model artifacts, MLflow experiment tracking and a notebook MLflow UI manager.

set_mlflow_tracking_uri

def set_mlflow_tracking_uri(
    uri:NoneType=None
):

Point MLflow at a tracking store for fastMONAI runs (opt-in).

Call once, before mlflow.set_experiment/mlflow.start_run or before creating fastMONAI’s MLflow callbacks, to choose where runs are logged.

Args: uri: A tracking URI (e.g. "sqlite:///runs.db" or an HTTP server). If None, uses fastMONAI’s default sqlite store at the package/repo root.

Returns: The tracking URI that was set.

source

store_variables

def store_variables(
    config_fn:str | pathlib.Path, size:list, apply_reorder:bool, target_spacing:int | list
):

Save inference variables as JSON (numpy values coerced to native types).

Written as JSON for safe, pickle-free sharing; load_variables reads JSON and still falls back to legacy pickle files.

source

load_variables

def load_variables(
    config_fn:str | pathlib.Path
):

Load stored inference variables.

Tries JSON first. Falls back to legacy pickle ONLY for files named .pkl (so a non-JSON payload disguised as .json is refused rather than unpickled, since pickle can execute arbitrary code). Existing .pkl artifacts still load.

Args: config_fn: File path to load.

Returns: The deserialized [size, apply_reorder, target_spacing] list.

Patch-based inference settings

source

decode_model_output

def decode_model_output(
    logits:Tensor, output_spec:Mapping
):

Apply the declared safe output contract to raw model logits.

source

standard_inference_config

def standard_inference_config(
    apply_reorder:bool, target_spacing, item_tfms:list, output_spec:Mapping
)->dict:

Build a complete standard/whole-volume inference contract from transform specs.

source

patch_inference_config

def patch_inference_config(
    config, output_spec:Mapping
)->dict:

Build the complete versioned inference contract for a patch workflow.

For binary_segmentation outputs, config.binary_threshold must equal the output threshold, so patch inference and output decoding agree.

source

make_output_spec

def make_output_spec(
    kind:str, threshold:float=None, classes:NoneType=None, output_shape:NoneType=None, class_axis:int=1
)->dict:

Build a versioned, allow-listed inference output/decoding contract.

Supported kind values are binary_segmentation, multiclass_segmentation, classification, and regression.

source

patch_config_to_dict

def patch_config_to_dict(
    config, inference_only:bool=False
)->dict:

Return PatchConfig values as a JSON-native dictionary.

config may be a PatchConfig-like object or a mapping. Missing fields in historical mappings receive the current PatchConfig defaults. Set inference_only=True to omit training-only sampling, queue, and data-state settings.

Colab GPU check

source

print_colab_gpu_info

def print_colab_gpu_info():

Check if we have a GPU attached to the runtime.

torch.compile helpers

source

strip_compile_prefix

def strip_compile_prefix(
    state_dict, prefix:str='_orig_mod.'
):

Remove the torch.compile key prefix from a state dict, in place.

learn.save() on a compiled model writes keys prefixed with _orig_mod., which an uncompiled module cannot load.

Args: state_dict: A state dict, e.g. from torch.load. fastai checkpoints written with with_opt=True are {‘model’: …, ‘opt’: …}; pass the ‘model’ value. prefix: Key prefix to remove.

Returns: The same dict, so it can be passed straight to load_state_dict. A dict without the prefix is returned unchanged.

source

unwrap_compiled_model

def unwrap_compiled_model(
    model
):

Return model with any torch.compile wrapper removed.

torch.compile(model) returns an OptimizedModule that holds the original module as _orig_mod alongside TorchDynamo state bound to the torch build that created it. A learner exported while compiled loads on any torch, but raises on the first forward pass under a different one. The returned module shares its weights with the wrapper.

Args: model: An nn.Module, or a fastai Learner whose .model is unwrapped in place.

Returns: The innermost nn.Module, or the same Learner that was passed in. A model that was never compiled is returned unchanged.

Note: Handles torch.compile(model). model.compile() compiles in place and leaves no wrapper to remove.

Safetensors model artifacts

Models are saved as Safetensors weights with metadata that holds the model specification and the inference contract. Loading rebuilds the model only through registered, trusted constructors and wrappers.

source

load_model_artifact

def load_model_artifact(
    path:str | pathlib.Path, device:str='cpu'
):

Load a Safetensors model artifact.

source

load_safetensors_model

def load_safetensors_model(
    path:str | pathlib.Path, device:str='cpu'
):

Strict-load a registered Safetensors model and return it in eval mode.

source

save_safetensors_model

def save_safetensors_model(
    model, path:str | pathlib.Path, model_spec:Mapping, inference_config:Mapping, artifact_role:str,
    mlflow_run:str=None, extra_metadata:Mapping=None
)->Path:

Save a rebuildable nn.Module with the model-level Safetensors API.

source

read_safetensors_metadata

def read_safetensors_metadata(
    path:str | pathlib.Path
)->dict:

Read and validate parsed fastMONAI Safetensors metadata.

source

build_model_from_spec

def build_model_from_spec(
    spec:Mapping
):

Rebuild a model only through registered, trusted constructors and wrappers.

source

make_model_spec

def make_model_spec(
    arch_id:str, arch_kwargs:Mapping, wrapper_spec:NoneType=None
)->dict:

Create a validated, versioned model specification.

Wrappers are listed from inner to outer and use {'wrapper_id': str, 'wrapper_kwargs': dict}.

source

register_model_wrapper

def register_model_wrapper(
    wrapper_id:str, constructor:Callable, replace:bool=False
):

Register a trusted constructor(model, **kwargs) wrapper.

source

register_model_constructor

def register_model_constructor(
    arch_id:str, constructor:Callable, replace:bool=False
):

Register a trusted model constructor for Safetensors reconstruction.

MLflow experiment tracking

create_mlflow_callback builds a ModelTrackingCallback from a Learner. The callback logs hyperparameters, per-epoch metrics, the actual train/validation split and Safetensors model artifacts to MLflow.

source

ModelTrackingCallback

def ModelTrackingCallback(
    model_name:str, loss_function:str, item_tfms:list, size:list, target_spacing:list, apply_reorder:bool,
    experiment_name:str=None, run_name:str=None, auto_start:bool=False, patch_config:dict=None, model_spec:dict=None,
    output_spec:dict=None, inference_config:dict=None, extra_params:dict=None, extra_tags:dict=None,
    dataset_version:str=None, preprocessing_manifest:str | pathlib.Path=None, sample_id_col:str | int=None,
    log_split:bool=True
):

A FastAI callback for comprehensive MLflow experiment tracking.

This callback automatically logs hyperparameters, metrics, model artifacts, and configuration to MLflow during training. If a checkpoint callback (SaveModelCallback, EMACheckpoint, or any TrackerCallback with fname) is present, the best model checkpoint will also be logged as an artifact.

Supports auto-managed runs when created via create_mlflow_callback().

source

create_mlflow_callback

def create_mlflow_callback(
    learn, experiment_name:str=None, run_name:str=None, auto_start:bool=True, model_name:str=None,
    model_spec:dict=None, output_spec:dict=None, inference_config:dict=None, extra_params:dict=None,
    extra_tags:dict=None, dataset_version:str=None, preprocessing_manifest:str | pathlib.Path=None,
    sample_id_col:str | int=None, log_split:bool=True
)->ModelTrackingCallback:

Create MLflow tracking callback with auto-extracted configuration.

This factory function automatically extracts configuration from the Learner, eliminating the need to manually specify parameters like size, transforms, loss function, etc.

Auto-extracts from Learner: - Preprocessing: apply_reorder, target_spacing, size/patch_size - Transforms: item_tfms or pre_patch_tfms - Training: loss_func, model architecture

Args: learn: fastai Learner instance experiment_name: MLflow experiment name. If None, uses model name. run_name: MLflow run name. If None, auto-generates with timestamp. auto_start: If True, auto-starts/stops MLflow run in before_fit/after_fit. model_name: Override the auto-extracted model name (used as the experiment name when experiment_name is None). model_spec: Required allow-listed architecture reconstruction specification. output_spec: Required patch-workflow output/decoding contract. inference_config: Required complete config for standard workflows; includes output. extra_params: Additional parameters to log (e.g., {‘dropout’: 0.5}). extra_tags: MLflow tags to set on the run. dataset_version: Optional annotation/dataset fingerprint for tracking. preprocessing_manifest: Optional manifest returned by preprocess_dataset(). Its dataset/cache identities are validated, tagged, and logged. sample_id_col: Optional stable identity column for split_version. By default, the raw image path is used when available. log_split: If True, hashes and logs actual fastai train/val membership.

Returns: ModelTrackingCallback ready to use with learn.fit()

Example: >>> # Standard workflow: pass the complete inference contract. >>> model_spec = make_model_spec(“monai.unet_inner”, unet_kwargs) >>> inference_config = standard_inference_config( … apply_reorder=True, target_spacing=target_spacing, … item_tfms=inference_transforms_to_specs(item_tfms), … output_spec=make_output_spec(“binary_segmentation”, threshold=0.5)) >>> callback = create_mlflow_callback( … learn, experiment_name=“Task02_Heart”, … model_spec=model_spec, inference_config=inference_config) >>> learn.fit_flat_cos(200, lr, cbs=[callback, save_best]) >>> >>> # Patch workflow: patch settings come from the DataLoaders; pass output_spec. >>> model_spec = make_model_spec(“monai.unet”, unet_kwargs) >>> output_spec = make_output_spec(“multiclass_segmentation”, classes=2) >>> callback = create_mlflow_callback( … learn, experiment_name=“Task02_Heart”, … model_spec=model_spec, output_spec=output_spec) >>> learn.fit_one_cycle(30, lr, cbs=[callback, save_best])

Locating Safetensors model artifacts

source

find_model_artifacts

def find_model_artifacts(
    run_ids, artifact_role, expected_members:NoneType=None
):

Download explicitly selected Safetensors model artifacts from MLflow.

run_ids must map each stable member identifier to one immutable MLflow run ID. Experiment-name or latest-run discovery is intentionally unsupported: callers must declare exactly which trained models inference will use.

MLflow UI

MLflowUIManager starts a local mlflow ui server from a notebook and stops it with the kernel.

source

MLflowUIManager

def MLflowUIManager():

Launch and manage a local mlflow ui server from a notebook.

The UI’s lifetime is tied to the kernel: it is reaped when the interpreter exits, gracefully (via an atexit handler) or on a hard kill (via Linux PR_SET_PDEATHSIG, see [_die_with_parent](https://fastmonai.no/utils.html#_die_with_parent)), so closing/restarting the notebook never leaves an orphaned server holding the port. Call stop() to shut it down sooner. An externally-started UI is reused, not killed, and calling start_ui() again while its own server runs just shows the link.