garmentiq.utils.checkpoint

 1"""Checkpoint loading with a mismatch guard.
 2
 3Weights are loaded leniently so that partial checkpoints keep working, but a checkpoint
 4that does not match the model class would otherwise load nothing and leave a silently
 5random model. This module warns when few or none of the expected tensors are found.
 6"""
 7import warnings
 8
 9
10def load_state_dict_checked(model, state_dict, model_path: str = ""):
11    """
12    Loads weights into a model and warns when the checkpoint does not match it.
13
14    GarmentIQ loads every checkpoint with `strict=False`, because legitimate checkpoints
15    routinely carry extra tensors (a bundled tracker, an optimiser state, a frozen
16    backbone) or omit a head that is re-initialised. The cost of that leniency is that a
17    checkpoint meant for a different architecture loads *nothing* and silently yields a
18    randomly initialised model, which then produces plausible-looking but meaningless
19    output.
20
21    This helper keeps the lenient behaviour but makes that failure visible: if no tensor
22    in the checkpoint matches the model, or almost none do, a warning is raised naming the
23    file. It never blocks a load, so existing code keeps working.
24
25    Args:
26        model (torch.nn.Module): The model to load weights into.
27        state_dict (dict): The checkpoint tensors, already stripped of any prefixes.
28        model_path (str, optional): Path to the checkpoint, used in the warning message.
29                                    Default is "".
30
31    Returns:
32        torch.nn.modules.module._IncompatibleKeys: The result of `load_state_dict`, listing
33                                                   missing and unexpected keys.
34    """
35    expected = set(model.state_dict().keys())
36    provided = set(state_dict.keys())
37    matched = expected & provided
38
39    where = f" from {model_path!r}" if model_path else ""
40
41    if expected and not matched:
42        warnings.warn(
43            f"No weights{where} matched {type(model).__name__}: none of the "
44            f"{len(provided)} tensor(s) in the checkpoint correspond to the "
45            f"{len(expected)} the model expects. The model is left randomly initialised "
46            f"and its output will be meaningless. Check that the checkpoint matches the "
47            f"model class.",
48            RuntimeWarning,
49            stacklevel=3,
50        )
51    elif expected and len(matched) < 0.5 * len(expected):
52        warnings.warn(
53            f"Only {len(matched)} of {len(expected)} weights expected by "
54            f"{type(model).__name__} were found{where}. The remaining "
55            f"{len(expected) - len(matched)} are randomly initialised, which usually "
56            f"means the checkpoint does not match the model class.",
57            RuntimeWarning,
58            stacklevel=3,
59        )
60
61    return model.load_state_dict(state_dict, strict=False)
62
63
64__all__ = ["load_state_dict_checked"]