Fix potential bug in DefaultOptimizerConstructor (#355)

* fix: fix custom_keys whenweight_decay is not specified

* docs: small fix

* feat: add type check

* refactor: move if outside for

mmcv/runner/optimizer/default_constructor.py

@@ -87,11 +87,15 @@ class DefaultOptimizerConstructor:
             raise TypeError('paramwise_cfg should be None or a dict, '
                             f'but got {type(self.paramwise_cfg)}')
 
-        if ('custom_keys' in self.paramwise_cfg
-                and not isinstance(self.paramwise_cfg['custom_keys'], dict)):
-            raise TypeError(
-                'If specified, custom_keys must be a dict, '
-                f'but got {type(self.paramwise_cfg["custom_keys"])}')
+        if 'custom_keys' in self.paramwise_cfg:
+            if not isinstance(self.paramwise_cfg['custom_keys'], dict):
+                raise TypeError(
+                    'If specified, custom_keys must be a dict, '
+                    f'but got {type(self.paramwise_cfg["custom_keys"])}')
+            if self.base_wd is None:
+                for key in self.paramwise_cfg['custom_keys']:
+                    if 'decay_mult' in self.paramwise_cfg['custom_keys'][key]:
+                        raise ValueError('base_wd should not be None')
 
         # get base lr and weight decay
         # weight_decay must be explicitly specified if mult is specified
@@ -154,11 +158,11 @@ class DefaultOptimizerConstructor:
             for key in sorted_keys:
                 if key in f'{prefix}.{name}':
                     is_custom = True
-                    param_group['lr'] = self.base_lr * custom_keys[key].get(
-                        'lr_mult', 1.)
-                    param_group[
-                        'weight_decay'] = self.base_wd * custom_keys[key].get(
-                            'decay_mult', 1.)
+                    lr_mult = custom_keys[key].get('lr_mult', 1.)
+                    param_group['lr'] = self.base_lr * lr_mult
+                    if self.base_wd is not None:
+                        decay_mult = custom_keys[key].get('decay_mult', 1.)
+                        param_group['weight_decay'] = self.base_wd * decay_mult
                     break
             if not is_custom:
                 # bias_lr_mult affects all bias parameters except for norm.bias

tests/test_optimizer.py

@@ -355,7 +355,7 @@ def test_default_optimizer_constructor():
     assert len(optimizer.param_groups) == len(model_parameters) == 11
     check_optimizer(optimizer, model, **paramwise_cfg)
 
-    # test DefaultOptimizerConstructor with custom_groups and ExampleModel
+    # test DefaultOptimizerConstructor with custom_keys and ExampleModel
     model = ExampleModel()
     optimizer_cfg = dict(
         type='SGD', lr=base_lr, weight_decay=base_wd, momentum=momentum)
@@ -372,7 +372,16 @@ def test_default_optimizer_constructor():
         # custom_keys should be a dict
         paramwise_cfg_ = dict(custom_keys=[0.1, 0.0001])
         optim_constructor = DefaultOptimizerConstructor(
-            optim_constructor, paramwise_cfg_)
+            optimizer_cfg, paramwise_cfg_)
+        optimizer = optim_constructor(model)
+
+    with pytest.raises(ValueError):
+        # if 'decay_mult' is specified in custom_keys, weight_decay should be
+        # specified
+        optimizer_cfg_ = dict(type='SGD', lr=0.01)
+        paramwise_cfg_ = dict(custom_keys={'.backbone': dict(decay_mult=0.5)})
+        optim_constructor = DefaultOptimizerConstructor(
+            optimizer_cfg_, paramwise_cfg_)
         optimizer = optim_constructor(model)
 
     optim_constructor = DefaultOptimizerConstructor(optimizer_cfg,
@@ -435,6 +444,52 @@ def test_default_optimizer_constructor():
                     assert param_groups[i][setting] == settings[
                         setting], f'{name} {setting}'
 
+    # test DefaultOptimizerConstructor with custom_keys and ExampleModel 2
+    model = ExampleModel()
+    optimizer_cfg = dict(type='SGD', lr=base_lr, momentum=momentum)
+    paramwise_cfg = dict(custom_keys={'param1': dict(lr_mult=10)})
+
+    optim_constructor = DefaultOptimizerConstructor(optimizer_cfg,
+                                                    paramwise_cfg)
+    optimizer = optim_constructor(model)
+    # check optimizer type and default config
+    assert isinstance(optimizer, torch.optim.SGD)
+    assert optimizer.defaults['lr'] == base_lr
+    assert optimizer.defaults['momentum'] == momentum
+    assert optimizer.defaults['weight_decay'] == 0
+
+    # check params groups
+    param_groups = optimizer.param_groups
+
+    groups = []
+    group_settings = []
+    # group 1, matches of 'param1'
+    groups.append(['param1', 'sub.param1'])
+    group_settings.append({
+        'lr': base_lr * 10,
+        'momentum': momentum,
+        'weight_decay': 0,
+    })
+    # group 2, default group
+    groups.append([
+        'sub.conv1.weight', 'sub.conv1.bias', 'sub.gn.weight', 'sub.gn.bias',
+        'conv1.weight', 'conv2.weight', 'conv2.bias', 'bn.weight', 'bn.bias'
+    ])
+    group_settings.append({
+        'lr': base_lr,
+        'momentum': momentum,
+        'weight_decay': 0
+    })
+
+    assert len(param_groups) == 11
+    for i, (name, param) in enumerate(model.named_parameters()):
+        assert torch.equal(param_groups[i]['params'][0], param)
+        for group, settings in zip(groups, group_settings):
+            if name in group:
+                for setting in settings:
+                    assert param_groups[i][setting] == settings[
+                        setting], f'{name} {setting}'
+
 
 def test_torch_optimizers():
     torch_optimizers = [