(model, filename, strict=False, revise_keys=[(r"^module\.", "")])
| 61 | |
| 62 | |
| 63 | def load_pretrain(model, filename, strict=False, revise_keys=[(r"^module\.", "")]): |
| 64 | checkpoint = torch.load(filename) |
| 65 | # OrderedDict is a subclass of dict |
| 66 | if not isinstance(checkpoint, dict): |
| 67 | raise RuntimeError(f"No state_dict found in checkpoint file {filename}") |
| 68 | # get state_dict from checkpoint |
| 69 | if "state_dict" in checkpoint: |
| 70 | state_dict = checkpoint["state_dict"] |
| 71 | elif "model" in checkpoint: |
| 72 | state_dict = checkpoint["model"] |
| 73 | else: |
| 74 | state_dict = checkpoint |
| 75 | # strip prefix of state_dict |
| 76 | for p, r in revise_keys: |
| 77 | state_dict = {re.sub(p, r, k): v for k, v in state_dict.items()} |
| 78 | # load state_dict |
| 79 | load_state_dict(model, state_dict, strict) |
| 80 | return checkpoint |
nothing calls this directly
no test coverage detected