Support regular expression in the mapping arg of copy_model_state (#…

…6917)

Part of #6552.

### Description
After PR #6835, we have added `copy_model_args` in the `load` API which
can help us update the state_dict flexibly.

https://github.com/KumoLiu/MONAI/blob/93a149a611b66153cf804b31a7b36a939e2e593a/monai/bundle/scripts.py#L397

Given this [issue](#6552),
we need to be able to filter the model's weights flexibly.
In `copy_model_state`, we already have a "mapping" arg, the filter will
be more flexible if we can support regular expression in the mapping.
This PR mainly added the support for regular expression for "mapping"
arg.

In the
[example](#6552 (comment))
in this [issue](#6552),
after this PR, we can do something like:
```
exclude_vars = "encoder.mask_token|encoder.norm.weight|encoder.norm.bias|out.conv.conv.weight|out.conv.conv.bias"
mapping={"encoder.layers(.*).0.0.": "swinViT.layers(.*).0."}
dst_dict, updated_keys, unchanged_keys = copy_model_state(
       model, ssl_weights, exclude_vars=exclude_vars, mapping=mapping
)
```

Additionally, based on the comments of Eric
[here](#6552 (comment)),
I totally agree, we could add a handler to make the pipeline easier to
implement, but perhaps this task is no need to set as a "BundleTodo" for
MONAIv1.3 but as an enhancement for MONAI near future.
What do you think? @ericspod @wyli @Nic-Ma 

### Types of changes
<!--- Put an `x` in all the boxes that apply, and remove the not
applicable items -->
- [x] Non-breaking change (fix or new feature that would not break
existing functionality).
- [ ] Breaking change (fix or new feature that would cause existing
functionality to change).
- [ ] New tests added to cover the changes.
- [ ] Integration tests passed locally by running `./runtests.sh -f -u
--net --coverage`.
- [ ] Quick tests passed locally by running `./runtests.sh --quick
--unittests --disttests`.
- [ ] In-line docstrings updated.
- [ ] Documentation updated, tested `make html` command in the `docs/`
folder.

---------

Signed-off-by: KumoLiu <[email protected]>

monai/networks/nets/swin_unetr.py

@@ -1055,3 +1055,52 @@ def forward(self, x, normalize=True):
         x4 = self.layers4[0](x3.contiguous())
         x4_out = self.proj_out(x4, normalize)
         return [x0_out, x1_out, x2_out, x3_out, x4_out]
+
+
+def filter_swinunetr(key, value):
+    """
+    A filter function used to filter the pretrained weights from [1], then the weights can be loaded into MONAI SwinUNETR Model.
+    This function is typically used with `monai.networks.copy_model_state`
+    [1] "Valanarasu JM et al., Disruptive Autoencoders: Leveraging Low-level features for 3D Medical Image Pre-training
+    <https://arxiv.org/abs/2307.16896>"
+
+    Args:
+        key: the key in the source state dict used for the update.
+        value: the value in the source state dict used for the update.
+
+    Examples::
+
+        import torch
+        from monai.apps import download_url
+        from monai.networks.utils import copy_model_state
+        from monai.networks.nets.swin_unetr import SwinUNETR, filter_swinunetr
+
+        model = SwinUNETR(img_size=(96, 96, 96), in_channels=1, out_channels=3, feature_size=48)
+        resource = (
+            "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/ssl_pretrained_weights.pth"
+        )
+        ssl_weights_path = "./ssl_pretrained_weights.pth"
+        download_url(resource, ssl_weights_path)
+        ssl_weights = torch.load(ssl_weights_path)["model"]
+
+        dst_dict, loaded, not_loaded = copy_model_state(model, ssl_weights, filter_func=filter_swinunetr)
+
+    """
+    if key in [
+        "encoder.mask_token",
+        "encoder.norm.weight",
+        "encoder.norm.bias",
+        "out.conv.conv.weight",
+        "out.conv.conv.bias",
+    ]:
+        return None
+
+    if key[:8] == "encoder.":
+        if key[8:19] == "patch_embed":
+            new_key = "swinViT." + key[8:]
+        else:
+            new_key = "swinViT." + key[8:18] + key[20:]
+
+        return new_key, value
+    else:
+        return None

monai/networks/utils.py

@@ -478,6 +478,7 @@ def copy_model_state(
     mapping=None,
     exclude_vars=None,
     inplace=True,
+    filter_func=None,
 ):
     """
     Compute a module state_dict, of which the keys are the same as `dst`. The values of `dst` are overwritten
@@ -490,7 +491,7 @@ def copy_model_state(
 
     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.
+        src: a pytorch module or state dict 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]`
@@ -499,6 +500,8 @@ def copy_model_state(
             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`.
+        filter_func: a filter function used to filter the weights to be loaded.
+            See 'filter_swinunetr' in "monai.networks.nets.swin_unetr.py".
 
     Examples:
         .. code-block:: python
@@ -536,6 +539,12 @@ def copy_model_state(
                 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))

tests/test_swin_unetr.py

@@ -11,15 +11,20 @@
 
 from __future__ import annotations
 
+import os
+import tempfile
 import unittest
 from unittest import skipUnless
 
 import torch
 from parameterized import parameterized
 
+from monai.apps import download_url
 from monai.networks import eval_mode
-from monai.networks.nets.swin_unetr import PatchMerging, PatchMergingV2, SwinUNETR
+from monai.networks.nets.swin_unetr import PatchMerging, PatchMergingV2, SwinUNETR, filter_swinunetr
+from monai.networks.utils import copy_model_state
 from monai.utils import optional_import
+from tests.utils import assert_allclose, skip_if_downloading_fails, skip_if_quick, testing_data_config
 
 einops, has_einops = optional_import("einops")
 
@@ -51,6 +56,14 @@
                             case_idx += 1
                             TEST_CASE_SWIN_UNETR.append(test_case)
 
+TEST_CASE_FILTER = [
+    [
+        {"img_size": (96, 96, 96), "in_channels": 1, "out_channels": 14, "feature_size": 48, "use_checkpoint": True},
+        "swinViT.layers1.0.blocks.0.norm1.weight",
+        torch.tensor([0.9473, 0.9343, 0.8566, 0.8487, 0.8065, 0.7779, 0.6333, 0.5555]),
+    ]
+]
+
 
 class TestSWINUNETR(unittest.TestCase):
     @parameterized.expand(TEST_CASE_SWIN_UNETR)
@@ -93,6 +106,24 @@ def test_patch_merging(self):
         t = PatchMerging(dim)(torch.zeros((1, 21, 20, 20, dim)))
         self.assertEqual(t.shape, torch.Size([1, 11, 10, 10, 20]))
 
+    @parameterized.expand(TEST_CASE_FILTER)
+    @skip_if_quick
+    def test_filter_swinunetr(self, input_param, key, value):
+        with skip_if_downloading_fails():
+            with tempfile.TemporaryDirectory() as tempdir:
+                file_name = "ssl_pretrained_weights.pth"
+                data_spec = testing_data_config("models", f"{file_name.split('.', 1)[0]}")
+                weight_path = os.path.join(tempdir, file_name)
+                download_url(
+                    data_spec["url"], weight_path, hash_val=data_spec["hash_val"], hash_type=data_spec["hash_type"]
+                )
+
+                ssl_weight = torch.load(weight_path)["model"]
+                net = SwinUNETR(**input_param)
+                dst_dict, loaded, not_loaded = copy_model_state(net, ssl_weight, filter_func=filter_swinunetr)
+                assert_allclose(dst_dict[key][:8], value, atol=1e-4, rtol=1e-4, type_test=False)
+                self.assertTrue(len(loaded) == 157 and len(not_loaded) == 2)
+
 
 if __name__ == "__main__":
     unittest.main()

tests/testing_data/data_config.json

@@ -128,6 +128,11 @@
             "url": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/se_resnext50_32x4d-a260b3a4.pth",
             "hash_type": "sha256",
             "hash_val": "a260b3a40f82dfe37c58d26a612bcf7bef0d27c6fed096226b0e4e9fb364168e"
+        },
+        "ssl_pretrained_weights": {
+            "url": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/ssl_pretrained_weights.pth",
+            "hash_type": "sha256",
+            "hash_val": "c3564f40a6a051d3753a6d8fae5cc8eaf21ce8d82a9a3baf80748d15664055e8"
         }
     },
     "configs": {