Function load_task_model_checkpoint

modelscope/utils/checkpoint.py:159–571

Load model checkpoint file and feed the parameters into the model. Args: model_to_load: The model to be load model_local_dir: The actual checkpoint dir on local disk. default_dtype: Set the default float type by 'torch.set_default_dtype' load_state_fn: An opt

(model_to_load,
                               model_local_dir,
                               default_dtype=None,
                               load_state_fn=None,
                               **kwargs)

def load_task_model_checkpoint(model_to_load,
                               model_local_dir,
                               default_dtype=None,
                               load_state_fn=None,
                               **kwargs):
    """
    Load model checkpoint file and feed the parameters into the model.
    Args:
        model_to_load: The model to be load
        model_local_dir: The actual checkpoint dir on local disk.
        default_dtype: Set the default float type by 'torch.set_default_dtype'
        load_state_fn: An optional load_state_fn used to load state_dict into the model.

    Returns:

    """

    def _add_head_prefix_to_state_dict(state_dicts, head_prefix,
                                       expected_keys_without_head_prefix,
                                       missing_keys):
        new_state_dict = OrderedDict()
        for name, module in state_dicts.items():
            if name in expected_keys_without_head_prefix:
                name_with_head = '.'.join([head_prefix, name])
                new_state_dict[name_with_head] = module
                expected_keys_without_head_prefix.remove(name)
                missing_keys = list(set(missing_keys) - set([name_with_head]))
            else:
                new_state_dict[name] = module

        missing_head_keys = []
        if len(expected_keys_without_head_prefix) > 0:
            missing_head_keys = expected_keys_without_head_prefix.copy()
        return new_state_dict, missing_head_keys, missing_keys

    def _find_mismatched_keys(
        state_dicts,
        model_state_dict,
        loaded_keys,
        prefix,
        add_prefix_to_model,
        remove_prefix_from_model,
        ignore_mismatched_sizes,
    ):
        mismatched_key = []
        if ignore_mismatched_sizes:
            for checkpoint_key in loaded_keys:
                model_key = checkpoint_key
                if remove_prefix_from_model:
                    # The model key starts with `prefix` but `checkpoint_key` doesn't, so we add it.
                    model_key = f'{prefix}.{checkpoint_key}'
                elif add_prefix_to_model:
                    # The model key doesn't start with `prefix` but `checkpoint_key` does, so we remove it.
                    model_key = '.'.join(checkpoint_key.split('.')[1:])

                if model_key in model_state_dict:
                    model_shape = model_state_dict[model_key].shape
                    checkpoint_shape = state_dicts[checkpoint_key].shape