Args: model (Module): Module to load checkpoint. filename (str): Accept local filepath, URL, ``torchvision://xxx``, ``open-mmlab://xxx``. Please refer to ``docs/model_zoo.md`` for details. map_location (str): Same as :func:`torch.load`. st
(model, state_dict, load_pos_embed, strict=False, pos_embed_interp=True, logger=None)
| 666 | |
| 667 | |
| 668 | def load_checkpoint(model, state_dict, load_pos_embed, strict=False, pos_embed_interp=True, logger=None): |
| 669 | """ |
| 670 | Args: |
| 671 | model (Module): Module to load checkpoint. |
| 672 | filename (str): Accept local filepath, URL, ``torchvision://xxx``, |
| 673 | ``open-mmlab://xxx``. Please refer to ``docs/model_zoo.md`` for |
| 674 | details. |
| 675 | map_location (str): Same as :func:`torch.load`. |
| 676 | strict (bool): Whether to allow different params for the model and |
| 677 | checkpoint. |
| 678 | logger (:mod:`logging.Logger` or None): The logger for error message. |
| 679 | |
| 680 | Returns: |
| 681 | dict or OrderedDict: The loaded checkpoint. |
| 682 | """ |
| 683 | # checkpoint = _load_checkpoint(filename, map_location) |
| 684 | # OrderedDict is a subclass of dict |
| 685 | # if not isinstance(checkpoint, dict): |
| 686 | # raise RuntimeError( |
| 687 | # f'No state_dict found in checkpoint file {filename}') |
| 688 | # get state_dict from checkpoint |
| 689 | if 'pos_embed' in state_dict: |
| 690 | if load_pos_embed: |
| 691 | if not pos_embed_interp: |
| 692 | pass |
| 693 | elif model.use_cls_token: |
| 694 | state_dict['pos_embed'] = interpolate_pos_embed_with_cls_token(pos_embed_checkpoint=state_dict['pos_embed'], |
| 695 | patch_shape=model.patch_embed.patch_shape, |
| 696 | num_extra_tokens=1) |
| 697 | else: |
| 698 | state_dict['pos_embed'] = interpolate_pos_embed(pos_embed_checkpoint=state_dict['pos_embed'], |
| 699 | patch_shape=model.patch_embed.patch_shape, |
| 700 | num_extra_tokens=1) |
| 701 | else: |
| 702 | del state_dict['pos_embed'] |
| 703 | print("checkpoint pos_embed removed") |
| 704 | model_dict = model.state_dict() |
| 705 | load_dict = { |
| 706 | k: v for k, v in state_dict.items() if k in model_dict.keys() |
| 707 | } |
| 708 | print("Missing keys: {}".format(list(set(model_dict) - set(load_dict)))) |
| 709 | |
| 710 | load_state_dict(model, state_dict, strict, logger) |
| 711 | |
| 712 | |
| 713 | def load_state_dict(module, state_dict, strict=False, logger=None): |
no test coverage detected