Repository: Project-MONAI/MONAI

Issue:
Utilities for partially loading weights from a state_dict
**Is your feature request related to a problem? Please describe.**
A utility function is needed to set a model's state_dict partially from another one.
this is required by https://github.com/Project-MONAI/MONAI/issues/2273 and https://github.com/Project-MONAI/MONAI/discussions/2329

an implementation is already part of the checkpoint loader:
https://github.com/Project-MONAI/MONAI/blob/a5364621712fe6b3de584223a4e28be638bc0fc4/monai/handlers/checkpoint_loader.py#L95-L106



Discussion:
Yes, we can extract this logic to be a utility function.

Thanks.


Resolving patch:
diff --git a/monai/handlers/checkpoint_loader.py b/monai/handlers/checkpoint_loader.py
--- a/monai/handlers/checkpoint_loader.py
+++ b/monai/handlers/checkpoint_loader.py
@@ -13,8 +13,8 @@
 from typing import TYPE_CHECKING, Dict, Optional
 
 import torch
-import torch.nn as nn
 
+from monai.networks.utils import copy_model_state
 from monai.utils import exact_version, optional_import
 
 Events, _ = optional_import("ignite.engine", "0.4.4", exact_version, "Events")
@@ -99,11 +99,7 @@ def __call__(self, engine: Engine) -> None:
 
             # skip items that don't match data shape
             for k, obj in self.load_dict.items():
-                if isinstance(obj, (nn.DataParallel, nn.parallel.DistributedDataParallel)):
-                    obj = obj.module
-                if isinstance(obj, torch.nn.Module):
-                    d = obj.state_dict()
-                    checkpoint[k] = {k: v for k, v in checkpoint[k].items() if k in d and v.shape == d[k].shape}
+                checkpoint[k] = copy_model_state(obj, checkpoint, inplace=False)[0]
 
         # save current max epochs setting in the engine, don't overwrite it if larger than max_epochs in checkpoint
         prior_max_epochs = engine.state.max_epochs
diff --git a/monai/networks/__init__.py b/monai/networks/__init__.py
--- a/monai/networks/__init__.py
+++ b/monai/networks/__init__.py
@@ -10,6 +10,7 @@
 # limitations under the License.
 
 from .utils import (
+    copy_model_state,
     eval_mode,
     icnr_init,
     normal_init,
diff --git a/monai/networks/utils.py b/monai/networks/utils.py
--- a/monai/networks/utils.py
+++ b/monai/networks/utils.py
@@ -11,10 +11,11 @@
 """
 Utilities and types for defining networks, these depend on PyTorch.
 """
-
+import re
 import warnings
+from collections import OrderedDict
 from contextlib import contextmanager
-from typing import Any, Callable, Optional, Sequence
+from typing import Any, Callable, Mapping, Optional, Sequence, Union
 
 import torch
 import torch.nn as nn
@@ -30,6 +31,7 @@
     "pixelshuffle",
     "eval_mode",
     "train_mode",
+    "copy_model_state",
 ]
 
 
@@ -310,3 +312,82 @@ def train_mode(*nets: nn.Module):
         # Return required networks to eval_list
         for n in eval_list:
             n.eval()
+
+
+def copy_model_state(
+    dst: Union[torch.nn.Module, Mapping],
+    src: Union[torch.nn.Module, Mapping],
+    dst_prefix="",
+    mapping=None,
+    exclude_vars=None,
+    inplace=True,
+):
+    """
+    Compute a module state_dict, of which the keys are the same as `dst`. The values of `dst` are overwritten
+    by the ones from `src` whenever their keys match. The method provides additional `dst_prefix` for
+    the `dst` key when matching them. `mapping` can be a `{"src_key": "dst_key"}` dict, indicating
+    `dst[dst_prefix + dst_key] = src[src_key]`.
+    This function is mainly to return a model state dict
+    for loading the `src` model state into the `dst` model, `src` and `dst` can have different dict keys, but
+    their corresponding values normally have the same shape.
+
+    Args:
+        dst: a pytorch module or state dict to be updated.
+        src: a pytorch module or state dist used to get the values used for the update.
+        dst_prefix: `dst` key prefix, so that `dst[dst_prefix + src_key]`
+            will be assigned to the value of `src[src_key]`.
+        mapping: a `{"src_key": "dst_key"}` dict, indicating that `dst[dst_prefix + dst_key]`
+            to be assigned to the value of `src[src_key]`.
+        exclude_vars: a regular expression to match the `dst` variable names,
+            so that their values are not overwritten by `src`.
+        inplace: whether to set the `dst` module with the updated `state_dict` via `load_state_dict`.
+            This option is only available when `dst` is a `torch.nn.Module`.
+
+    Examples:
+        .. code-block:: python
+
+            from monai.networks.nets import BasicUNet
+            from monai.networks.utils import copy_model_state
+
+            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.
+    """
+
+    if isinstance(src, (nn.DataParallel, nn.parallel.DistributedDataParallel)):
+        src = src.module
+    if isinstance(dst, (nn.DataParallel, nn.parallel.DistributedDataParallel)):
+        dst = dst.module
+    src_dict = src.state_dict() if isinstance(src, torch.nn.Module) else src
+    dst_dict = dst.state_dict() if isinstance(dst, torch.nn.Module) else dst
+    dst_dict = OrderedDict(dst_dict)
+
+    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)
+
+    updated_keys = sorted(set(updated_keys))
+    unchanged_keys = sorted(set(all_keys).difference(updated_keys))
+    print(f"'dst' model updated: {len(updated_keys)} of {len(dst_dict)} variables.")
+    if inplace and isinstance(dst, torch.nn.Module):
+        dst.load_state_dict(dst_dict)
+    return dst_dict, updated_keys, unchanged_keys


Test patch:
diff --git a/tests/test_copy_model_state.py b/tests/test_copy_model_state.py
new file mode 100644
--- /dev/null
+++ b/tests/test_copy_model_state.py
@@ -0,0 +1,182 @@
+# Copyright 2020 - 2021 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+import torch
+from parameterized import parameterized
+
+from monai.networks.utils import copy_model_state
+from monai.utils import set_determinism
+
+
+class _TestModelOne(torch.nn.Module):
+    def __init__(self, n_n, n_m, n_class):
+        super(_TestModelOne, self).__init__()
+        self.layer = torch.nn.Linear(n_n, n_m)
+        self.class_layer = torch.nn.Linear(n_m, n_class)
+
+    def forward(self, x):
+        x = self.layer(x)
+        x = self.class_layer(x)
+        return x
+
+
+class _TestModelTwo(torch.nn.Module):
+    def __init__(self, n_n, n_m, n_d, n_class):
+        super(_TestModelTwo, self).__init__()
+        self.layer = torch.nn.Linear(n_n, n_m)
+        self.layer_1 = torch.nn.Linear(n_m, n_d)
+        self.class_layer = torch.nn.Linear(n_d, n_class)
+
+    def forward(self, x):
+        x = self.layer(x)
+        x = self.layer_1(x)
+        x = self.class_layer(x)
+        return x
+
+
+TEST_CASES = []
+__devices = ("cpu", "cuda") if torch.cuda.is_available() else ("cpu",)
+for _x in __devices:
+    for _y in __devices:
+        TEST_CASES.append((_x, _y))
+
+
+class TestModuleState(unittest.TestCase):
+    def tearDown(self):
+        set_determinism(None)
+
+    @parameterized.expand(TEST_CASES)
+    def test_set_state(self, device_0, device_1):
+        set_determinism(0)
+        model_one = _TestModelOne(10, 20, 3)
+        model_two = _TestModelTwo(10, 20, 10, 4)
+        model_one.to(device_0)
+        model_two.to(device_1)
+        model_dict, ch, unch = copy_model_state(model_one, model_two)
+        x = np.random.randn(4, 10)
+        x = torch.tensor(x, device=device_0, dtype=torch.float32)
+        output = model_one(x).detach().cpu().numpy()
+        expected = np.array(
+            [
+                [-0.36076584, -0.03177825, -0.7702266],
+                [-0.0526831, -0.15855855, -0.01149344],
+                [-0.3760508, -0.22485238, -0.0634037],
+                [0.5977675, -0.67991066, 0.1919502],
+            ]
+        )
+        np.testing.assert_allclose(output, expected, atol=1e-3)
+        self.assertEqual(len(ch), 2)
+        self.assertEqual(len(unch), 2)
+
+    @parameterized.expand(TEST_CASES)
+    def test_set_full_state(self, device_0, device_1):
+        set_determinism(0)
+        model_one = _TestModelOne(10, 20, 3)
+        model_two = _TestModelOne(10, 20, 3)
+        model_one.to(device_0)
+        model_two.to(device_1)
+        # test module input
+        model_dict, ch, unch = copy_model_state(model_one, model_two)
+        # test dict input
+        model_dict, ch, unch = copy_model_state(model_dict, model_two)
+        x = np.random.randn(4, 10)
+        x = torch.tensor(x, device=device_0, dtype=torch.float32)
+        output = model_one(x).detach().cpu().numpy()
+        model_two.to(device_0)
+        output_1 = model_two(x).detach().cpu().numpy()
+        np.testing.assert_allclose(output, output_1, atol=1e-3)
+        self.assertEqual(len(ch), 4)
+        self.assertEqual(len(unch), 0)
+
+    @parameterized.expand(TEST_CASES)
+    def test_set_exclude_vars(self, device_0, device_1):
+        set_determinism(0)
+        model_one = _TestModelOne(10, 20, 3)
+        model_two = _TestModelTwo(10, 20, 10, 4)
+        model_one.to(device_0)
+        model_two.to(device_1)
+        # test skip layer.bias
+        model_dict, ch, unch = copy_model_state(model_one, model_two, exclude_vars="layer.bias")
+        x = np.random.randn(4, 10)
+        x = torch.tensor(x, device=device_0, dtype=torch.float32)
+        output = model_one(x).detach().cpu().numpy()
+        expected = np.array(
+            [
+                [-0.34172416, 0.0375042, -0.98340976],
+                [-0.03364138, -0.08927619, -0.2246768],
+                [-0.35700908, -0.15556987, -0.27658707],
+                [0.61680925, -0.6106281, -0.02123314],
+            ]
+        )
+        np.testing.assert_allclose(output, expected, atol=1e-3)
+        self.assertEqual(len(ch), 1)
+        self.assertEqual(len(unch), 3)
+
+    @parameterized.expand(TEST_CASES)
+    def test_set_map_across(self, device_0, device_1):
+        set_determinism(0)
+        model_one = _TestModelOne(10, 10, 3)
+        model_two = _TestModelTwo(10, 10, 10, 4)
+        model_one.to(device_0)
+        model_two.to(device_1)
+        # test weight map
+        model_dict, ch, unch = copy_model_state(
+            model_one, model_two, mapping={"layer_1.weight": "layer.weight", "layer_1.bias": "layer_1.weight"}
+        )
+        model_one.load_state_dict(model_dict)
+        x = np.random.randn(4, 10)
+        x = torch.tensor(x, device=device_0, dtype=torch.float32)
+        output = model_one(x).detach().cpu().numpy()
+        expected = np.array(
+            [
+                [0.8244487, -0.19650555, 0.65723234],
+                [0.71239626, 0.25617486, 0.5247122],
+                [0.24168758, 1.0301148, 0.39089814],
+                [0.25791705, 0.8653245, 0.14833644],
+            ]
+        )
+        np.testing.assert_allclose(output, expected, atol=1e-3)
+        self.assertEqual(len(ch), 2)
+        self.assertEqual(len(unch), 2)
+
+    @parameterized.expand(TEST_CASES)
+    def test_set_prefix(self, device_0, device_1):
+        set_determinism(0)
+        model_one = torch.nn.Sequential(_TestModelOne(10, 20, 3))
+        model_two = _TestModelTwo(10, 20, 10, 4)
+        model_one.to(device_0)
+        model_two.to(device_1)
+        # test skip layer.bias
+        model_dict, ch, unch = copy_model_state(
+            model_one, model_two, dst_prefix="0.", exclude_vars="layer.bias", inplace=False
+        )
+        model_one.load_state_dict(model_dict)
+        x = np.random.randn(4, 10)
+        x = torch.tensor(x, device=device_0, dtype=torch.float32)
+        output = model_one(x).detach().cpu().numpy()
+        expected = np.array(
+            [
+                [-0.360766, -0.031778, -0.770227],
+                [-0.052683, -0.158559, -0.011493],
+                [-0.376051, -0.224852, -0.063404],
+                [0.597767, -0.679911, 0.19195],
+            ]
+        )
+        np.testing.assert_allclose(output, expected, atol=1e-3)
+        self.assertEqual(len(ch), 2)
+        self.assertEqual(len(unch), 2)
+
+
+if __name__ == "__main__":
+    unittest.main()