Function _load_checkpoint

modelscope/utils/checkpoint.py:278–485

(
        model,
        state_dict,
        load_state_fn,
        ignore_mismatched_sizes,
        _fast_init,
    )

        return error_msgs

    def _load_checkpoint(
        model,
        state_dict,
        load_state_fn,
        ignore_mismatched_sizes,
        _fast_init,
    ):
        # Retrieve missing & unexpected_keys
        model_state_dict = model.state_dict()
        expected_keys = list(model_state_dict.keys())
        keys_from_pretrained = list(state_dict.keys())

        prefix = model.base_model_prefix

        # during loading stage, base model prefix is complicated, should consider remove or add
        if len(prefix) > 0:
            # nlp: encoder, decoder
            pretrained_has_prefix_module = any(
                s.startswith(prefix) for s in keys_from_pretrained)
            model_expects_prefix_module = any(
                s.startswith(prefix) for s in expected_keys)
        else:
            # nlp:encoder-decoder, cv:backbone-head,
            pretrained_has_prefix_module = False
            model_expects_prefix_module = False

        remove_prefix_from_model = not pretrained_has_prefix_module and model_expects_prefix_module
        add_prefix_to_model = pretrained_has_prefix_module and not model_expects_prefix_module

        if remove_prefix_from_model:
            expected_keys_not_base_model_prefixed = [
                s for s in expected_keys if not s.startswith(prefix)
            ]
            expected_keys = [
                '.'.join(s.split('.')[1:]) if s.startswith(prefix) else s
                for s in expected_keys
            ]
        elif add_prefix_to_model:
            # backbone only
            expected_keys = ['.'.join([prefix, s]) for s in expected_keys]
            expected_keys_not_base_model_prefixed = []

        missing_keys = list(set(expected_keys) - set(keys_from_pretrained))
        unexpected_keys = list(set(keys_from_pretrained) - set(expected_keys))

        # during loading stage head prefix is simple, add or not add
        prefix_heads = model.head_prefix
        expected_head_keys_without_head_prefix = []
        missing_head_keys = []
        unexpected_head_keys = []
        pretrained_has_prefix_head = dict()
        head_prefix_keys = dict()

        # only for case of head mismatched with state-dict
        if len(prefix_heads) > 0 and len(unexpected_keys) > 0:
            if isinstance(prefix_heads, str):
                prefix_heads = [prefix_heads]