scxpand.util.model_loading

Unified model loading utilities for scXpand.

This module provides robust, consistent model loading functionality across all model types to prevent silent failures and ensure proper error handling.

Functions

load_and_validate_model(model, model_dir, device)	Load state dict into model with validation.
load_model_parameters(model_dir, param_class)	Load model parameters from JSON file.
load_model_state_dict(model_dir, device[, ...])	Load model state dictionary from training checkpoint.
validate_model_directory(model_dir, ...)	Validate that model directory contains required files.

Exceptions

ModelLoadingError

Custom exception for model loading failures.

exception scxpand.util.model_loading.ModelLoadingError

Custom exception for model loading failures.

scxpand.util.model_loading.load_and_validate_model(model, model_dir, device, model_name='model')

Load state dict into model with validation.

Parameters:

• model (Module) – Initialized model instance

• model_dir (Path) – Directory containing model files

• device (device | str) – Device to load model on

• model_name (str (default: 'model')) – Name of the model (for logging)

Return type:

Module

Returns:

Model with loaded weights in eval mode

Raises:

• ModelLoadingError – If loading fails

• RuntimeError – If state dict is incompatible with model

scxpand.util.model_loading.load_model_parameters(model_dir, param_class, param_file='parameters.json')

Load model parameters from JSON file.

Parameters:

• model_dir (Path) – Directory containing parameter file

• param_class (type) – Parameter class to instantiate

• param_file (str (default: 'parameters.json')) – Name of parameter file

Return type:

Any

Returns:

Instantiated parameter object

Raises:

• FileNotFoundError – If parameter file doesn’t exist

• json.JSONDecodeError – If parameter file is invalid JSON

• TypeError – If parameter instantiation fails

scxpand.util.model_loading.load_model_state_dict(model_dir, device, model_name='model')

Load model state dictionary from training checkpoint.

This function loads model state dictionary from best_ckpt.pt only. No fallback to legacy files - ensures consistent, reliable model loading.

Parameters:

• model_dir (Path) – Directory containing model files

• device (device | str) – Device to load model on

• model_name (str (default: 'model')) – Name of the model (for logging)

Return type:

dict[str, Any]

Returns:

Model state dictionary

Raises:

• ModelLoadingError – If no valid model checkpoint is found

• FileNotFoundError – If model directory doesn’t exist

• RuntimeError – If model loading fails due to corruption or incompatibility

scxpand.util.model_loading.validate_model_directory(model_dir, required_files)

Validate that model directory contains required files.

Parameters:

• model_dir (Path) – Directory to validate

• required_files (list[str]) – List of required file names

Raises:

• FileNotFoundError – If directory or required files don’t exist

• ValueError – If path is not a directory

Return type:

None