Repository: Project-MONAI/MONAI

Issue:
Freeze layers for transfer learning
Is it possible to freeze layers in the UNet?
I want to freeze some layers in order to do transfer learning.

However, when I use model.children, the whole network is printed. I am not able to get specific layers.

So, is this possible, or do you recommand something else for transfer learning?

Thank you. 

Best,
Kirsten


Discussion:

Resolving patch:
diff --git a/monai/optimizers/__init__.py b/monai/optimizers/__init__.py
--- a/monai/optimizers/__init__.py
+++ b/monai/optimizers/__init__.py
@@ -10,3 +10,4 @@
 # limitations under the License.
 
 from .novograd import Novograd
+from .utils import *
diff --git a/monai/optimizers/utils.py b/monai/optimizers/utils.py
new file mode 100644
--- /dev/null
+++ b/monai/optimizers/utils.py
@@ -0,0 +1,76 @@
+# Copyright 2020 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.
+
+from typing import Callable, Sequence
+
+import torch
+
+from monai.utils import ensure_tuple
+
+
+def generate_param_groups(
+    network: torch.nn.Module,
+    layer_matches: Sequence[Callable],
+    match_types: Sequence[str],
+    lr_values: Sequence[float],
+    include_others: bool = True,
+):
+    """
+    Utility function to generate parameter groups with different LR values for optimizer.
+    The output parameter groups have the same order as `layer_match` functions.
+
+    Args:
+        network: source network to generate parameter groups from.
+        layer_matches: a list of callable functions to select or filter out network layer groups,
+            for "select" type, the input will be the `network`, for "filter" type,
+            the input will be every item of `network.named_parameters()`.
+        match_types: a list of tags to identify the matching type corresponding to the `layer_matches` functions,
+            can be "select" or "filter".
+        lr_values: a list of LR values corresponding to the `layer_matches` functions.
+        include_others: whether to incude the rest layers as the last group, default to True.
+
+    It's mainly used to set different init LR values for different network elements, for example::
+
+        net = Unet(dimensions=3, in_channels=1, out_channels=3, channels=[2, 2, 2], strides=[1, 1, 1])
+        print(net)  # print out network components to select expected items
+        print(net.named_parameters())  # print out all the named parameters to filter out expected items
+        params = generate_param_groups(
+            network=net,
+            layer_matches=[lambda x: x.model[-1], lambda x: "conv.weight" in x],
+            match_types=["select", "filter"],
+            lr_values=[1e-2, 1e-3],
+        )
+        optimizer = torch.optim.Adam(params, 1e-4)
+
+    """
+    layer_matches = ensure_tuple(layer_matches)
+    match_types = ensure_tuple(match_types)
+    lr_values = ensure_tuple(lr_values)
+    if len(layer_matches) != len(lr_values) or len(layer_matches) != len(match_types):
+        raise ValueError("length of layer_match callable functions, match types and LR values should be the same.")
+
+    params = list()
+    _layers = list()
+    for func, ty, lr in zip(layer_matches, match_types, lr_values):
+        if ty == "select":
+            layer_params = func(network).parameters()
+        elif ty == "filter":
+            layer_params = filter(func, network.named_parameters())
+        else:
+            raise ValueError(f"unsuppoted layer match type: {ty}.")
+
+        params.append({"params": layer_params, "lr": lr})
+        _layers.extend(list(map(id, layer_params)))
+
+    if include_others:
+        params.append({"params": filter(lambda p: id(p) not in _layers, network.parameters())})
+
+    return params


Test patch:
diff --git a/tests/test_generate_param_groups.py b/tests/test_generate_param_groups.py
new file mode 100644
--- /dev/null
+++ b/tests/test_generate_param_groups.py
@@ -0,0 +1,89 @@
+# Copyright 2020 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 torch
+from parameterized import parameterized
+
+from monai.networks.nets import Unet
+from monai.optimizers import generate_param_groups
+from monai.utils import ensure_tuple
+
+TEST_CASE_1 = [
+    {
+        "layer_matches": [lambda x: x.model[-1]],
+        "match_types": ["select"],
+        "lr_values": [1],
+    },
+    (1, 100),
+]
+
+TEST_CASE_2 = [
+    {
+        "layer_matches": [lambda x: x.model[-1], lambda x: x.model[-2], lambda x: x.model[-3]],
+        "match_types": ["select", "select", "select"],
+        "lr_values": [1, 2, 3],
+    },
+    (1, 2, 3, 100),
+]
+
+TEST_CASE_3 = [
+    {
+        "layer_matches": [lambda x: x.model[2][1].conv[0].conv],
+        "match_types": ["select"],
+        "lr_values": [1],
+    },
+    (1, 100),
+]
+
+TEST_CASE_4 = [
+    {
+        "layer_matches": [lambda x: x.model[-1], lambda x: "conv.weight" in x],
+        "match_types": ["select", "filter"],
+        "lr_values": [1, 2],
+    },
+    (1, 2, 100),
+]
+
+TEST_CASE_5 = [
+    {
+        "layer_matches": [lambda x: x.model[-1]],
+        "match_types": ["select"],
+        "lr_values": [1],
+        "include_others": False,
+    },
+    (1),
+]
+
+
+class TestGenerateParamGroups(unittest.TestCase):
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5])
+    def test_lr_values(self, input_param, expected_values):
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        net = Unet(
+            dimensions=3,
+            in_channels=1,
+            out_channels=3,
+            channels=(16, 32, 64),
+            strides=(2, 2),
+            num_res_units=1,
+        ).to(device)
+
+        params = generate_param_groups(network=net, **input_param)
+        optimizer = torch.optim.Adam(params, 100)
+
+        for param_group, value in zip(optimizer.param_groups, ensure_tuple(expected_values)):
+            torch.testing.assert_allclose(param_group["lr"], value)
+
+
+if __name__ == "__main__":
+    unittest.main()