Skip to content

← nn.Module Mechanics step 5 of 7

Medium Primitives

The weights that did not load

model.load_state_dict(checkpoint)

By default this is strict: every key in the checkpoint must exist in the model and vice versa, or it raises. That default is doing you a favour, and the interesting failures come from turning it off.

model.load_state_dict(checkpoint, strict=False)

returns a named tuple instead of raising:

result = model.load_state_dict(checkpoint, strict=False)
result.missing_keys       # in the model, absent from the checkpoint
result.unexpected_keys    # in the checkpoint, absent from the model

strict=False is genuinely necessary when loading a pretrained backbone into a model with a new head, or when a checkpoint predates a layer you added. The mistake is not using it; the mistake is not reading the return value.

A silently empty load looks exactly like a successful one. The model runs, it produces plausible outputs, and every weight is whatever __init__ left there.

Where the keys go wrong

  • DataParallel prefixes. Saving from a wrapped model puts module. in front of every key. Loading into an unwrapped one makes every key unexpected and every key missing at the same time.
  • A renamed attribute. Renaming self.fc to self.head renames every key under it.
  • Buffers with persistent=False. Deliberately absent from the checkpoint, and therefore permanently in missing_keys. Worth knowing so you do not chase it.

The habit

result = model.load_state_dict(ckpt, strict=False)
assert not result.unexpected_keys, result.unexpected_keys
print("randomly initialised:", result.missing_keys)

Unexpected keys almost always mean a mismatch worth stopping for. Missing keys are sometimes intentional, so print them and read them.

Your task

def load_report(checkpoint_keys: list) -> dict

Build a model with parameters a.weight, a.bias and b.weight, and try to load a checkpoint containing only the named keys. Return missing and unexpected, both sorted, without raising.