mapping_networks.checkpoint.load_checkpoint

mapping_networks.checkpoint.load_checkpoint(path: str | Path, model: MappingModel, *, trainer: MappingTrainer | None = None, map_location: Any = None, strict_latents: bool = True) → CheckpointSchema

Load a checkpoint and restore model (and optionally trainer) state.

The target’s ParameterSpec is validated against the checkpoint before any state is applied. If the architecture differs a CheckpointCompatibilityError is raised and the model is left unchanged.

Parameters:

  • path – Path to the .pt checkpoint file.

  • model – A MappingModel whose generator state will be restored.

  • trainer – Optional MappingTrainer whose optimizer, scheduler, epoch, and GradScaler state will be restored when present in the checkpoint.

  • map_location – Passed to torch.load for device remapping.

  • strict_latents – When True, load_state_dict is called in strict mode — every saved latent and buffer key must match the generator. Set to False to allow partial loading (e.g. subsets of layers).

Returns:

The parsed CheckpointSchema so callers can inspect metadata, config, or trainer_state independently.

Raises:

  • FileNotFoundError – If the checkpoint file does not exist.

  • CheckpointCompatibilityError – If schema version or architecture is incompatible.

Example:

schema = load_checkpoint("checkpoints/epoch10.pt", model=model, trainer=trainer)
print(f"Resumed from epoch {schema.trainer_state['epoch']}")