feat: print optimizer param groups when verbose (#1149)

Prints which module parameters match each optimizer param group when verbose is set, both at init and on set_params. Fixes #291.

Author: aiedwardyi

skorch/net.py

@@ -41,6 +41,7 @@
 from skorch.exceptions import SkorchAttributeError
 from skorch.exceptions import SkorchTrainingImpossibleError
 from skorch.history import History
+from skorch.setter import format_param_group_msg
 from skorch.setter import optimizer_setter
 from skorch.utils import _TorchLoadUnpickler
 from skorch.utils import _identity
@@ -2030,7 +2031,13 @@ def _get_params_for_optimizer(self, prefix, named_parameters):
             matches = [i for i, (name, _) in enumerate(params) if
                        fnmatch.fnmatch(name, pattern)]
             if matches:
-                p = [params.pop(i)[1] for i in reversed(matches)]
+                # pop high indices first so earlier indices stay valid
+                matched = [params.pop(i) for i in reversed(matches)]
+                p = [param for _, param in matched]
+                if self.verbose:
+                    # show names in the order they were matched
+                    matched_names = [name for name, _ in reversed(matched)]
+                    print(format_param_group_msg(group, matched_names))
                 pgroups.append({'params': p, **group})
 
         if params:

skorch/setter.py

@@ -1,6 +1,41 @@
 """Setter functions for virtual params such as ``optimizer__lr``."""
 import re
 
+# param names can be arbitrarily long, keep the verbose message bounded
+MAX_PARAM_GROUP_MSG_LEN = 200
+
+
+def format_param_group_msg(group_config, param_names):
+    """Message for which module params a param group config applies to."""
+    if not param_names:
+        msg = (
+            "Setting param group {} for parameters that are not among the "
+            "module's learnable parameters (this may be unintended).".format(
+                group_config,
+            )
+        )
+    else:
+        msg = "Setting param group {} for {}.".format(
+            group_config,
+            ', '.join(param_names),
+        )
+    if len(msg) > MAX_PARAM_GROUP_MSG_LEN:
+        msg = msg[:MAX_PARAM_GROUP_MSG_LEN - 3] + '...'
+    return msg
+
+
+def _param_names_for_tensors(net, tensors):
+    """Map optimizer tensors back to module parameter names when possible."""
+    tensor_ids = {id(t) for t in tensors}
+    names = []
+    get_params = getattr(net, 'get_all_learnable_params', None)
+    if get_params is None:
+        return names
+    for name, p in get_params():
+        if id(p) in tensor_ids:
+            names.append(name)
+    return names
+
 
 def _extract_optimizer_param_name_and_group(optimizer_name, param):
     """Extract param group and param name from the given parameter name.
@@ -44,6 +79,8 @@ def _set_optimizer_param(optimizer, param_group, param_name, value):
     for group in groups:
         group[param_name] = value
 
+    return groups
+
 
 def optimizer_setter(
         net, param, value, optimizer_attr='optimizer_', optimizer_name='optimizer'
@@ -62,9 +99,18 @@ def optimizer_setter(
         param_group, param_name = _extract_optimizer_param_name_and_group(
             optimizer_name, param)
 
-    _set_optimizer_param(
+    groups = _set_optimizer_param(
         optimizer=getattr(net, optimizer_attr),
         param_group=param_group,
         param_name=param_name,
         value=value
     )
+
+    # only report for a specific param group; a global set (e.g. optimizer__lr)
+    # touches every param and is not what #291 asks to surface
+    if getattr(net, 'verbose', 0) and param_group != 'all':
+        tensors = []
+        for group in groups:
+            tensors.extend(group.get('params', []))
+        param_names = _param_names_for_tensors(net, tensors)
+        print(format_param_group_msg({param_name: value}, param_names))

skorch/tests/test_net.py

@@ -1396,6 +1396,62 @@ def test_optimizer_param_groups(self, net_cls, module_cls):
         assert net.optimizer_.param_groups[1]['lr'] == 0.5
         assert net.optimizer_.param_groups[2]['lr'] == net.lr
 
+    @pytest.mark.parametrize('param_groups, expected_count, expected_msgs', [
+        ([], 0, []),
+        (
+            [('sequential.0.*', {'lr': 0.1})],
+            1,
+            [
+                "Setting param group {'lr': 0.1} for",
+                'sequential.0.weight',
+                'sequential.0.bias',
+            ],
+        ),
+        (
+            [
+                ('sequential.0.*', {'lr': 0.1}),
+                ('sequential.3.*', {'lr': 0.5}),
+            ],
+            2,
+            [
+                "Setting param group {'lr': 0.1} for",
+                "Setting param group {'lr': 0.5} for",
+                'sequential.0.weight',
+                'sequential.3.weight',
+            ],
+        ),
+    ])
+    def test_optimizer_param_groups_verbose_prints(
+            self, net_cls, module_cls, data, param_groups, expected_count,
+            expected_msgs, capsys):
+        # fit instead of initialize so accidental repetitions during
+        # training would be caught by the count
+        X, y = data
+        net = net_cls(
+            module_cls,
+            verbose=1,
+            max_epochs=2,
+            optimizer__param_groups=param_groups,
+        )
+        net.fit(X, y)
+        out = capsys.readouterr().out
+        assert out.count('Setting param group') == expected_count
+        for expected in expected_msgs:
+            assert expected in out
+
+    def test_optimizer_param_groups_silent_when_verbose_0(
+            self, net_cls, module_cls, capsys):
+        net = net_cls(
+            module_cls,
+            verbose=0,
+            optimizer__param_groups=[
+                ('sequential.0.*', {'lr': 0.1}),
+            ],
+        )
+        net.initialize()
+        out = capsys.readouterr().out
+        assert 'Setting param group' not in out
+
     def test_module_params_in_init(self, net_cls, module_cls, data):
         X, y = data
 

skorch/tests/test_setter.py

@@ -73,3 +73,40 @@ def test_only_specific_param_group_updated(self, setter, net_optim_dummy,
         assert updated_group_new[0][sub_param] == value
         assert all(old[sub_param] == new[sub_param] for old, new in zip(
             static_groups_pre, static_groups_new))
+
+    def test_set_params_verbose_prints_param_group(self, setter, capsys):
+        from skorch import NeuralNetClassifier
+        from skorch.toy import make_classifier
+
+        net = NeuralNetClassifier(make_classifier(), verbose=1, max_epochs=1)
+        net.initialize()
+        setter(net, 'optimizer__param_groups__0__lr', 0.03)
+        out = capsys.readouterr().out
+        assert "Setting param group {'lr': 0.03} for" in out
+        assert 'sequential.0.weight' in out
+
+    def test_set_params_verbose_silent_for_global_param(self, setter, capsys):
+        from skorch import NeuralNetClassifier
+        from skorch.toy import make_classifier
+
+        net = NeuralNetClassifier(make_classifier(), verbose=1, max_epochs=1)
+        net.initialize()
+        setter(net, 'optimizer__lr', 0.03)
+        out = capsys.readouterr().out
+        assert 'Setting param group' not in out
+
+    def test_format_param_group_msg_no_known_params(self):
+        from skorch.setter import format_param_group_msg
+
+        msg = format_param_group_msg({'lr': 0.1}, [])
+        assert "{'lr': 0.1}" in msg
+        assert 'this may be unintended' in msg
+
+    def test_format_param_group_msg_truncated(self):
+        from skorch.setter import MAX_PARAM_GROUP_MSG_LEN
+        from skorch.setter import format_param_group_msg
+
+        names = ['sequential.{}.weight'.format(i) for i in range(100)]
+        msg = format_param_group_msg({'lr': 0.1}, names)
+        assert len(msg) == MAX_PARAM_GROUP_MSG_LEN
+        assert msg.endswith('...')