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, logger=None)
| 567 | |
| 568 | |
| 569 | def load_checkpoint(model, state_dict, load_pos_embed, strict=False, logger=None): |
| 570 | """ |
| 571 | Args: |
| 572 | model (Module): Module to load checkpoint. |
| 573 | filename (str): Accept local filepath, URL, ``torchvision://xxx``, |
| 574 | ``open-mmlab://xxx``. Please refer to ``docs/model_zoo.md`` for |
| 575 | details. |
| 576 | map_location (str): Same as :func:`torch.load`. |
| 577 | strict (bool): Whether to allow different params for the model and |
| 578 | checkpoint. |
| 579 | logger (:mod:`logging.Logger` or None): The logger for error message. |
| 580 | |
| 581 | Returns: |
| 582 | dict or OrderedDict: The loaded checkpoint. |
| 583 | """ |
| 584 | # checkpoint = _load_checkpoint(filename, map_location) |
| 585 | # OrderedDict is a subclass of dict |
| 586 | # if not isinstance(checkpoint, dict): |
| 587 | # raise RuntimeError( |
| 588 | # f'No state_dict found in checkpoint file {filename}') |
| 589 | # get state_dict from checkpoint |
| 590 | if 'pos_embed' in state_dict: |
| 591 | if load_pos_embed: |
| 592 | state_dict['pos_embed'] = interpolate_pos_embed(pos_embed_checkpoint=state_dict['pos_embed'], |
| 593 | patch_shape=model.patch_embed.patch_shape, |
| 594 | num_extra_tokens=1) |
| 595 | else: |
| 596 | del state_dict['pos_embed'] |
| 597 | print("checkpoint pos_embed removed") |
| 598 | |
| 599 | model_dict = model.state_dict() |
| 600 | load_dict = { |
| 601 | k: v for k, v in state_dict.items() if k in model_dict.keys() |
| 602 | } |
| 603 | print("Missing keys: {}".format(list(set(model_dict) - set(load_dict)))) |
| 604 | |
| 605 | load_state_dict(model, state_dict, strict, logger) |
| 606 | return checkpoint |
| 607 | |
| 608 | |
| 609 | def load_state_dict(module, state_dict, strict=False, logger=None): |
no test coverage detected