Load state_dict to a module. This method is modified from :meth:`torch.nn.Module.load_state_dict`. Default value for ``strict`` is set to ``False`` and the message for param mismatch will NOT be shown if strict is False. Args: module (Module): Module that receives the state
(module, state_dict, strict=False, logger=None)
| 39 | |
| 40 | |
| 41 | def load_state_dict(module, state_dict, strict=False, logger=None): |
| 42 | """Load state_dict to a module. |
| 43 | |
| 44 | This method is modified from :meth:`torch.nn.Module.load_state_dict`. |
| 45 | Default value for ``strict`` is set to ``False`` and the message for |
| 46 | param mismatch will NOT be shown if strict is False. |
| 47 | |
| 48 | Args: |
| 49 | module (Module): Module that receives the state_dict. |
| 50 | state_dict (OrderedDict): Weights. |
| 51 | strict (bool): whether to strictly enforce that the keys |
| 52 | in :attr:`state_dict` match the keys returned by this module's |
| 53 | :meth:`~torch.nn.Module.state_dict` function. Default: ``False``. |
| 54 | logger (:obj:`logging.Logger`, optional): Logger to log the error |
| 55 | message. If not specified, print function will be used. |
| 56 | """ |
| 57 | unexpected_keys = [] |
| 58 | all_missing_keys = [] |
| 59 | err_msg = [] |
| 60 | |
| 61 | metadata = getattr(state_dict, '_metadata', None) |
| 62 | state_dict = state_dict.copy() |
| 63 | if metadata is not None: |
| 64 | state_dict._metadata = metadata |
| 65 | |
| 66 | # use _load_from_state_dict to enable checkpoint version control |
| 67 | def load(module, prefix=''): |
| 68 | # recursively check parallel module in case that the model has a |
| 69 | # complicated structure, e.g., nn.Module(nn.Module(DDP)) |
| 70 | if is_module_wrapper(module): |
| 71 | module = module.module |
| 72 | local_metadata = {} if metadata is None else metadata.get( |
| 73 | prefix[:-1], {}) |
| 74 | module._load_from_state_dict(state_dict, prefix, local_metadata, True, |
| 75 | all_missing_keys, unexpected_keys, |
| 76 | err_msg) |
| 77 | for name, child in module._modules.items(): |
| 78 | if child is not None: |
| 79 | load(child, prefix + name + '.') |
| 80 | |
| 81 | load(module) |
| 82 | load = None # break load->load reference cycle |
| 83 | |
| 84 | # ignore "num_batches_tracked" of BN layers |
| 85 | missing_keys = [ |
| 86 | key for key in all_missing_keys if 'num_batches_tracked' not in key |
| 87 | ] |
| 88 | |
| 89 | if unexpected_keys: |
| 90 | err_msg.append('unexpected key in source ' |
| 91 | f'state_dict: {", ".join(unexpected_keys)}\n') |
| 92 | if missing_keys: |
| 93 | err_msg.append( |
| 94 | f'missing keys in source state_dict: {", ".join(missing_keys)}\n') |
| 95 | |
| 96 | if strict: |
| 97 | rank, _ = get_dist_info() |
| 98 | if len(err_msg) > 0 and rank == 0: |
no test coverage detected