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 be shown even if strict is False. Args: module (Module): Module that receives the stat
(module, state_dict, strict=False, logger=None)
| 711 | |
| 712 | |
| 713 | def load_state_dict(module, state_dict, strict=False, logger=None): |
| 714 | """Load state_dict to a module. |
| 715 | |
| 716 | This method is modified from :meth:`torch.nn.Module.load_state_dict`. |
| 717 | Default value for ``strict`` is set to ``False`` and the message for |
| 718 | param mismatch will be shown even if strict is False. |
| 719 | |
| 720 | Args: |
| 721 | module (Module): Module that receives the state_dict. |
| 722 | state_dict (OrderedDict): Weights. |
| 723 | strict (bool): whether to strictly enforce that the keys |
| 724 | in :attr:`state_dict` match the keys returned by this module's |
| 725 | :meth:`~torch.nn.Module.state_dict` function. Default: ``False``. |
| 726 | logger (:obj:`logging.Logger`, optional): Logger to log the error |
| 727 | message. If not specified, print function will be used. |
| 728 | """ |
| 729 | unexpected_keys = [] |
| 730 | all_missing_keys = [] |
| 731 | err_msg = [] |
| 732 | |
| 733 | metadata = getattr(state_dict, '_metadata', None) |
| 734 | state_dict = state_dict.copy() |
| 735 | if metadata is not None: |
| 736 | state_dict._metadata = metadata |
| 737 | |
| 738 | # use _load_from_state_dict to enable checkpoint version control |
| 739 | def load(module, prefix=''): |
| 740 | # recursively check parallel module in case that the model has a |
| 741 | # complicated structure, e.g., nn.Module(nn.Module(DDP)) |
| 742 | # if is_module_wrapper(module): |
| 743 | # module = module.module |
| 744 | local_metadata = {} if metadata is None else metadata.get( |
| 745 | prefix[:-1], {}) |
| 746 | module._load_from_state_dict(state_dict, prefix, local_metadata, True, |
| 747 | all_missing_keys, unexpected_keys, |
| 748 | err_msg) |
| 749 | for name, child in module._modules.items(): |
| 750 | if child is not None: |
| 751 | load(child, prefix + name + '.') |
| 752 | |
| 753 | load(module) |
| 754 | load = None # break load->load reference cycle |
| 755 | |
| 756 | # ignore "num_batches_tracked" of BN layers |
| 757 | missing_keys = [ |
| 758 | key for key in all_missing_keys if 'num_batches_tracked' not in key |
| 759 | ] |
| 760 | |
| 761 | if unexpected_keys: |
| 762 | err_msg.append('unexpected key in source ' |
| 763 | f'state_dict: {", ".join(unexpected_keys)}\n') |
| 764 | if missing_keys: |
| 765 | err_msg.append( |
| 766 | f'missing keys in source state_dict: {", ".join(missing_keys)}\n') |
| 767 | |
| 768 | rank = dist.get_rank() |
| 769 | |
| 770 | if len(err_msg) > 0 and rank == 0: |
no test coverage detected