pimm.export.load_pretrained

load_pretrained(model: Module, checkpoint_path: str, prefix: str | None = None, remove_prefix: bool = True, key_mapping: Dict[str, str] | None = None, strict: bool = True, device: str | device = 'cpu', filter_fn: Callable[[Dict[str, Any], Module], Dict[str, Any]] | None = None) → Dict[str, Any] | None

Load pretrained weights into model with flexible key matching.

This function supports various loading scenarios:

• Full checkpoint loading (all keys match)

• Partial loading with prefix filtering (e.g., load only “backbone.” weights)

• Key remapping (e.g., map “student.backbone.” to “backbone.”)

• Custom filtering via filter_fn

Example

# Load full checkpoint
load_pretrained(model, "checkpoint.pth")

# Load only backbone weights from a Sonata checkpoint
load_pretrained(model, "sonata_checkpoint.pth", prefix="student.backbone.")

# Load backbone weights and remap keys
load_pretrained(
    model,
    "sonata_checkpoint.pth",
    prefix="student.backbone.",
    key_mapping={"backbone.": ""},  # remove the "backbone." prefix
)

Parameters:

• model – Model to load weights into.

• checkpoint_path – Path to checkpoint file.

• prefix – Optional prefix to filter keys (e.g., “backbone.”, “student.backbone.”). If provided, only keys starting with this prefix will be loaded.

• remove_prefix – If True and prefix is provided, remove prefix from keys.

• key_mapping – Optional dictionary to remap keys (old_key -> new_key).

• strict – Whether to strictly enforce all keys match.

• device – Device to load checkpoint to.

• filter_fn – Optional function(state_dict, model) -> filtered_state_dict for custom filtering logic.

Returns:

Incompatible keys (missing_keys, unexpected_keys) if strict=False, None otherwise.