(module, state_dict, strict=False, logger=None)
| 18 | |
| 19 | |
| 20 | def load_state_dict(module, state_dict, strict=False, logger=None): |
| 21 | unexpected_keys = [] |
| 22 | all_missing_keys = [] |
| 23 | err_msg = [] |
| 24 | |
| 25 | metadata = getattr(state_dict, "_metadata", None) |
| 26 | state_dict = state_dict.copy() |
| 27 | if metadata is not None: |
| 28 | state_dict._metadata = metadata |
| 29 | |
| 30 | # use _load_from_state_dict to enable checkpoint version control |
| 31 | def load(module, prefix=""): |
| 32 | # recursively check parallel module in case that the model has a |
| 33 | # complicated structure, e.g., nn.Module(nn.Module(DDP)) |
| 34 | local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {}) |
| 35 | module._load_from_state_dict( |
| 36 | state_dict, prefix, local_metadata, True, all_missing_keys, unexpected_keys, err_msg |
| 37 | ) |
| 38 | for name, child in module._modules.items(): |
| 39 | if child is not None: |
| 40 | load(child, prefix + name + ".") |
| 41 | |
| 42 | load(module) |
| 43 | load = None # break load->load reference cycle |
| 44 | |
| 45 | # ignore "num_batches_tracked" of BN layers |
| 46 | missing_keys = [key for key in all_missing_keys if "num_batches_tracked" not in key] |
| 47 | |
| 48 | if unexpected_keys: |
| 49 | err_msg.append(f"unexpected key in source state_dict: {', '.join(unexpected_keys)}\n") |
| 50 | if missing_keys: |
| 51 | err_msg.append(f"missing keys in source state_dict: {', '.join(missing_keys)}\n") |
| 52 | |
| 53 | rank, _ = get_dist_info() |
| 54 | if len(err_msg) > 0 and rank == 0: |
| 55 | err_msg.insert(0, "The model and loaded state dict do not match exactly\n") |
| 56 | err_msg = "\n".join(err_msg) |
| 57 | if strict: |
| 58 | raise RuntimeError(err_msg) |
| 59 | else: |
| 60 | print(err_msg) |
| 61 | |
| 62 | |
| 63 | def load_pretrain(model, filename, strict=False, revise_keys=[(r"^module\.", "")]): |
no test coverage detected