Repository: Project-MONAI/MONAI
File: monai/networks/utils.py (Python)

Code before:
            model_a_b, changed, unchanged = copy_model_state(
                model_a, model_b, exclude_vars="conv_0.conv_0", inplace=False)
            # dst model updated: 76 of 82 variables.
            model_a.load_state_dict(model_a_b)
            # <All keys matched successfully>

    Returns: an OrderedDict of the updated `dst` state, the changed, and unchanged keys.

    """
    src_dict = get_state_dict(src)
    dst_dict = OrderedDict(get_state_dict(dst))

    to_skip = {s_key for s_key in src_dict if exclude_vars and re.compile(exclude_vars).search(s_key)}

    # update dst with items from src
    all_keys, updated_keys = list(dst_dict), list()
    for s, val in src_dict.items():
        dst_key = f"{dst_prefix}{s}"
        if dst_key in dst_dict and dst_key not in to_skip and dst_dict[dst_key].shape == val.shape:
            dst_dict[dst_key] = val
            updated_keys.append(dst_key)
    for s in mapping if mapping else {}:
        dst_key = f"{dst_prefix}{mapping[s]}"
        src_keys = sorted({s_key for s_key in src_dict if s_key not in to_skip and re.compile(s).search(s_key)})
        dst_keys = sorted({d_key for d_key in dst_dict if d_key not in to_skip and re.compile(dst_key).search(d_key)})
        for _src_key, _dst_key in zip(src_keys, dst_keys):
            if dst_dict[_dst_key].shape != src_dict[_src_key].shape:
                warnings.warn(f"Param. shape changed from {dst_dict[_dst_key].shape} to {src_dict[_src_key].shape}.")
            dst_dict[_dst_key] = src_dict[_src_key]
            updated_keys.append(_dst_key)

    updated_keys = sorted(set(updated_keys))
    unchanged_keys = sorted(set(all_keys).difference(updated_keys))
    logger.info(f"'dst' model updated: {len(updated_keys)} of {len(dst_dict)} variables.")
    if inplace and isinstance(dst, torch.nn.Module):
        if isinstance(dst, (nn.DataParallel, nn.parallel.DistributedDataParallel)):
            dst = dst.module
        dst.load_state_dict(dst_dict)
    return dst_dict, updated_keys, unchanged_keys


def save_state(src: torch.nn.Module | dict, path: PathLike, **kwargs):
    """
    Save the state dict of input source data with PyTorch `save`.
    It can save `nn.Module`, `state_dict`, a dictionary of `nn.Module` or `state_dict`.
    And automatically convert the data parallel module to regular module.
    For example::

        save_state(net, path)
        save_state(net.state_dict(), path)
        save_state({"net": net, "opt": opt}, path)

Review comment (suggestion):
this `zip` assumes 1 to 1 mapping for the first N keys, I can see in practice there could be missing keys in the middle of `dst_keys` or `src_keys` and will cause mismatches... in my opinion defining this kind of search pattern is complicated, but if there's a use case I'm fine with this new feature.

Code after:

            model_a = BasicUNet(in_channels=1, out_channels=4)
            model_b = BasicUNet(in_channels=1, out_channels=2)
            model_a_b, changed, unchanged = copy_model_state(
                model_a, model_b, exclude_vars="conv_0.conv_0", inplace=False)
            # dst model updated: 76 of 82 variables.
            model_a.load_state_dict(model_a_b)
            # <All keys matched successfully>

    Returns: an OrderedDict of the updated `dst` state, the changed, and unchanged keys.

    """
    src_dict = get_state_dict(src)
    dst_dict = OrderedDict(get_state_dict(dst))

    to_skip = {s_key for s_key in src_dict if exclude_vars and re.compile(exclude_vars).search(s_key)}

    # update dst with items from src
    all_keys, updated_keys = list(dst_dict), list()
    for s, val in src_dict.items():
        dst_key = f"{dst_prefix}{s}"
        if dst_key in dst_dict and dst_key not in to_skip and dst_dict[dst_key].shape == val.shape:
            dst_dict[dst_key] = val
            updated_keys.append(dst_key)
    for s in mapping if mapping else {}:
        dst_key = f"{dst_prefix}{mapping[s]}"
        if dst_key in dst_dict and dst_key not in to_skip:
            if dst_dict[dst_key].shape != src_dict[s].shape:
                warnings.warn(f"Param. shape changed from {dst_dict[dst_key].shape} to {src_dict[s].shape}.")
            dst_dict[dst_key] = src_dict[s]
            updated_keys.append(dst_key)
    if filter_func is not None:
        for key, value in src_dict.items():
            new_pair = filter_func(key, value)
            if new_pair is not None and new_pair[0] not in to_skip:
                dst_dict[new_pair[0]] = new_pair[1]
                updated_keys.append(new_pair[0])

    updated_keys = sorted(set(updated_keys))
    unchanged_keys = sorted(set(all_keys).difference(updated_keys))
    logger.info(f"'dst' model updated: {len(updated_keys)} of {len(dst_dict)} variables.")
    if inplace and isinstance(dst, torch.nn.Module):
        if isinstance(dst, (nn.DataParallel, nn.parallel.DistributedDataParallel)):
            dst = dst.module
        dst.load_state_dict(dst_dict)
    return dst_dict, updated_keys, unchanged_keys


def save_state(src: torch.nn.Module | dict, path: PathLike, **kwargs):
    """
    Save the state dict of input source data with PyTorch `save`.