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)
| 607 | |
| 608 | |
| 609 | def load_state_dict(module, state_dict, strict=False, logger=None): |
| 610 | """Load state_dict to a module. |
| 611 | |
| 612 | This method is modified from :meth:`torch.nn.Module.load_state_dict`. |
| 613 | Default value for ``strict`` is set to ``False`` and the message for |
| 614 | param mismatch will be shown even if strict is False. |
| 615 | |
| 616 | Args: |
| 617 | module (Module): Module that receives the state_dict. |
| 618 | state_dict (OrderedDict): Weights. |
| 619 | strict (bool): whether to strictly enforce that the keys |
| 620 | in :attr:`state_dict` match the keys returned by this module's |
| 621 | :meth:`~torch.nn.Module.state_dict` function. Default: ``False``. |
| 622 | logger (:obj:`logging.Logger`, optional): Logger to log the error |
| 623 | message. If not specified, print function will be used. |
| 624 | """ |
| 625 | unexpected_keys = [] |
| 626 | all_missing_keys = [] |
| 627 | err_msg = [] |
| 628 | |
| 629 | metadata = getattr(state_dict, '_metadata', None) |
| 630 | state_dict = state_dict.copy() |
| 631 | if metadata is not None: |
| 632 | state_dict._metadata = metadata |
| 633 | |
| 634 | # use _load_from_state_dict to enable checkpoint version control |
| 635 | def load(module, prefix=''): |
| 636 | # recursively check parallel module in case that the model has a |
| 637 | # complicated structure, e.g., nn.Module(nn.Module(DDP)) |
| 638 | # if is_module_wrapper(module): |
| 639 | # module = module.module |
| 640 | local_metadata = {} if metadata is None else metadata.get( |
| 641 | prefix[:-1], {}) |
| 642 | module._load_from_state_dict(state_dict, prefix, local_metadata, True, |
| 643 | all_missing_keys, unexpected_keys, |
| 644 | err_msg) |
| 645 | for name, child in module._modules.items(): |
| 646 | if child is not None: |
| 647 | load(child, prefix + name + '.') |
| 648 | |
| 649 | load(module) |
| 650 | load = None # break load->load reference cycle |
| 651 | |
| 652 | # ignore "num_batches_tracked" of BN layers |
| 653 | missing_keys = [ |
| 654 | key for key in all_missing_keys if 'num_batches_tracked' not in key |
| 655 | ] |
| 656 | |
| 657 | if unexpected_keys: |
| 658 | err_msg.append('unexpected key in source ' |
| 659 | f'state_dict: {", ".join(unexpected_keys)}\n') |
| 660 | if missing_keys: |
| 661 | err_msg.append( |
| 662 | f'missing keys in source state_dict: {", ".join(missing_keys)}\n') |
| 663 | |
| 664 | rank = dist.get_rank() |
| 665 | |
| 666 | if len(err_msg) > 0 and rank == 0: |
no test coverage detected