CheckpointLoader

class CheckpointLoader(keywords='', replacement=None, replacements=None, strict=False)

Bases: HookBase

Load model weights and, when requested, resume optimizer/train state.

In before_train, before the trainer enters the loop, delegates to load_weight_and_resume(), which loads weights from cfg.weight (warm-start) or resumes a full run from cfg.resume. State-dict keys are rewritten by the configured replacement rules before loading, which is how a checkpoint trained under one module layout is mapped onto another (e.g. wrapping a backbone under module.model.backbone). Registered as CheckpointLoader.

Parameters:

• keywords (str) – Single substring to match in checkpoint keys when using the scalar form. Defaults to "" (matches everything / no-op).

• replacement (str, optional) – String to substitute for keywords in matched keys. When None, falls back to keywords (i.e. no rewrite). Defaults to None.

• replacements (dict, optional) – A {keyword: replacement} mapping to apply several key rewrites in a single load_state_dict call, so missing/unexpected keys are reported truthfully. Overlapping keywords are resolved by longest match, so dict order is irrelevant. Mutually exclusive in practice with the scalar form. Must be a dict or a TypeError is raised. Defaults to None.

• strict (bool) – Whether to require an exact key match when loading model weights (passed through to load_state_dict). Defaults to False.

Note

Warm-starting (cfg.weight) loads model weights only; full resume (cfg.resume) additionally restores optimizer, scheduler, and step state. After loading, judge success by inspecting the reported missing / unexpected keys rather than assuming silence means success.

Example

Add to cfg.hooks; once in before_train it loads cfg.weight (warm-start) or cfg.resume (full resume), rewriting state-dict keys by the given rules before load_state_dict:

hooks = [
    dict(type="CheckpointLoader", replacements={
        "module.backbone": "module.model.backbone",
        "module.decoder":  "module.model.decoder",
    }),
]
# → before the loop, loads cfg.weight/cfg.resume after renaming
#   "module.backbone.*" -> "module.model.backbone.*" (etc.); inspect the
#   reported missing/unexpected keys to confirm the load took

before_train()

Load configured weights before the trainer enters the train loop.