Skip to content

← nn.Module Mechanics step 5 of 7

Medium Primitives

Report missing and unexpected checkpoint keys

Jump to the task

A checkpoint can be missing weights your model expects, or contain names your model does not recognize. By default, load_state_dict raises for either kind of key mismatch. With strict=False, it reports those names instead:

import torch
import torch.nn as nn

model = nn.Linear(2, 2)
checkpoint = {"weight": torch.ones(2, 2), "old_head": torch.zeros(1)}
result = model.load_state_dict(checkpoint, strict=False)
print(result.missing_keys)     # ['bias']
print(result.unexpected_keys)  # ['old_head']

The weight was loaded. The missing bias kept its existing value, and the unrecognized old_head entry was not loaded into this model.

What strict=False does not ignore

It relaxes key matching, not every compatibility check. A recognized weight with the wrong shape still raises. For a new classification head with different dimensions, remove or transform the incompatible entries deliberately; simply setting strict=False does not make their shapes fit.

Check both returned lists against the differences you intended. An empty checkpoint can leave every parameter unchanged without raising under strict=False. Conversely, an empty mismatch report establishes key agreement, not that this is the correct checkpoint for your task.

Common causes of mismatched names

  • Saving a DataParallel wrapper's state_dict adds a module. prefix. The wrapped model's own state_dict, available through wrapper.module, omits it.
  • Renaming a registered layer from fc to head changes the prefix of its keys.
  • Adding or removing a layer changes the set of expected keys.

A buffer registered with persistent=False is not expected in the state_dict. Its absence is therefore not a missing key. If an old checkpoint includes that buffer, its name can instead be reported as unexpected.

Your task

def load_report(checkpoint_keys: list) -> dict

Use the supplied model with keys a.weight, a.bias and b.weight. Build a checkpoint with exactly the names supplied in checkpoint_keys. For recognized names, the starter copies correctly shaped values from another model. For unknown names, it uses a one-element zero tensor as a placeholder.

Load with strict=False and return the sorted lists missing and unexpected. The inputs contain distinct names, and recognized names always have valid shapes, so these cases should not raise.

For ["a.weight", "old.weight"], return {"missing": ["a.bias", "b.weight"], "unexpected": ["old.weight"]}. Do not filter unknown names out before loading: reporting them is part of the task.