File size: 2,872 Bytes
80cf062 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | from __future__ import annotations
from contextlib import contextmanager
from pathlib import Path
import sys
from typing import Any, Iterator, Mapping
import torch
from torch import nn
@contextmanager
def _temporary_module_path(source_root: str | Path | None) -> Iterator[None]:
if source_root is None:
yield
return
path = str(Path(source_root).resolve())
sys.path.insert(0, path)
try:
yield
finally:
try:
sys.path.remove(path)
except ValueError:
pass
def _extract_state_dict(checkpoint: Any) -> Mapping[str, torch.Tensor]:
if isinstance(checkpoint, Mapping):
for key in ("model", "model_state", "state_dict"):
candidate = checkpoint.get(key)
if isinstance(candidate, Mapping):
return candidate
if checkpoint and all(isinstance(key, str) for key in checkpoint):
return checkpoint
raise ValueError("Checkpoint does not contain a model state dictionary.")
def load_checkpoint(
model: nn.Module,
checkpoint_path: str | Path,
*,
strict: bool = True,
map_location: str | torch.device = "cpu",
source_root: str | Path | None = None,
) -> dict[str, Any]:
"""Load a trusted W-MAE checkpoint and return loading metadata.
``source_root`` is needed only for legacy checkpoints whose pickle payload
references modules from the original repository, such as ``utils.YParams``.
"""
checkpoint_path = Path(checkpoint_path)
if not checkpoint_path.is_file():
raise FileNotFoundError(f"Checkpoint does not exist: {checkpoint_path}")
with _temporary_module_path(source_root):
checkpoint = torch.load(checkpoint_path, map_location=map_location, weights_only=False)
state_dict = _extract_state_dict(checkpoint)
model_state = model.state_dict()
missing = sorted(set(model_state) - set(state_dict))
unexpected = sorted(set(state_dict) - set(model_state))
shape_mismatches = {
key: {"model": list(model_state[key].shape), "checkpoint": list(state_dict[key].shape)}
for key in model_state.keys() & state_dict.keys()
if tuple(model_state[key].shape) != tuple(state_dict[key].shape)
}
if strict and (missing or unexpected or shape_mismatches):
raise RuntimeError(
"Checkpoint is incompatible with the model: "
f"missing={missing}, unexpected={unexpected}, shape_mismatches={shape_mismatches}"
)
incompatible = model.load_state_dict(state_dict, strict=strict)
return {
"checkpoint_path": str(checkpoint_path),
"strict": strict,
"missing_keys": list(incompatible.missing_keys),
"unexpected_keys": list(incompatible.unexpected_keys),
"shape_mismatches": shape_mismatches,
"checkpoint_keys": len(state_dict),
}
|