test_optimizer.py 23.3 KB
Newer Older
1
import sys
2
import warnings
3
from unittest.mock import MagicMock
4
5
6
7
8
9
10
11

import pytest
import torch
import torch.nn as nn

from mmcv.runner import OPTIMIZER_BUILDERS, DefaultOptimizerConstructor
from mmcv.runner.optimizer import build_optimizer, build_optimizer_constructor
from mmcv.runner.optimizer.builder import TORCH_OPTIMIZERS
12
13
14
15
16
17
from mmcv.utils.ext_loader import check_ops_exist

OPS_AVAILABLE = check_ops_exist()
if not OPS_AVAILABLE:
    sys.modules['mmcv.ops'] = MagicMock(
        DeformConv2d=dict, ModulatedDeformConv2d=dict)
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40


class SubModel(nn.Module):

    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(2, 2, kernel_size=1, groups=2)
        self.gn = nn.GroupNorm(2, 2)
        self.param1 = nn.Parameter(torch.ones(1))

    def forward(self, x):
        return x


class ExampleModel(nn.Module):

    def __init__(self):
        super().__init__()
        self.param1 = nn.Parameter(torch.ones(1))
        self.conv1 = nn.Conv2d(3, 4, kernel_size=1, bias=False)
        self.conv2 = nn.Conv2d(4, 2, kernel_size=1)
        self.bn = nn.BatchNorm2d(2)
        self.sub = SubModel()
41
42
43
44
        if OPS_AVAILABLE:
            from mmcv.ops import DeformConv2dPack
            self.dcn = DeformConv2dPack(
                3, 4, kernel_size=3, deformable_groups=1)
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60

    def forward(self, x):
        return x


class ExampleDuplicateModel(nn.Module):

    def __init__(self):
        super().__init__()
        self.param1 = nn.Parameter(torch.ones(1))
        self.conv1 = nn.Sequential(nn.Conv2d(3, 4, kernel_size=1, bias=False))
        self.conv2 = nn.Sequential(nn.Conv2d(4, 2, kernel_size=1))
        self.bn = nn.BatchNorm2d(2)
        self.sub = SubModel()
        self.conv3 = nn.Sequential(nn.Conv2d(3, 4, kernel_size=1, bias=False))
        self.conv3[0] = self.conv1[0]
61
62
63
64
        if OPS_AVAILABLE:
            from mmcv.ops import DeformConv2dPack
            self.dcn = DeformConv2dPack(
                3, 4, kernel_size=3, deformable_groups=1)
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90

    def forward(self, x):
        return x


class PseudoDataParallel(nn.Module):

    def __init__(self):
        super().__init__()
        self.module = ExampleModel()

    def forward(self, x):
        return x


base_lr = 0.01
base_wd = 0.0001
momentum = 0.9


def check_default_optimizer(optimizer, model, prefix=''):
    assert isinstance(optimizer, torch.optim.SGD)
    assert optimizer.defaults['lr'] == base_lr
    assert optimizer.defaults['momentum'] == momentum
    assert optimizer.defaults['weight_decay'] == base_wd
    param_groups = optimizer.param_groups[0]
91
92
93
94
95
96
97
98
99
100
101
102
103
    if OPS_AVAILABLE:
        param_names = [
            'param1', 'conv1.weight', 'conv2.weight', 'conv2.bias',
            'bn.weight', 'bn.bias', 'sub.param1', 'sub.conv1.weight',
            'sub.conv1.bias', 'sub.gn.weight', 'sub.gn.bias', 'dcn.weight',
            'dcn.conv_offset.weight', 'dcn.conv_offset.bias'
        ]
    else:
        param_names = [
            'param1', 'conv1.weight', 'conv2.weight', 'conv2.bias',
            'bn.weight', 'bn.bias', 'sub.param1', 'sub.conv1.weight',
            'sub.conv1.bias', 'sub.gn.weight', 'sub.gn.bias'
        ]
104
105
106
107
108
109
110
    param_dict = dict(model.named_parameters())
    assert len(param_groups['params']) == len(param_names)
    for i in range(len(param_groups['params'])):
        assert torch.equal(param_groups['params'][i],
                           param_dict[prefix + param_names[i]])


111
112
113
114
115
116
117
118
119
def check_sgd_optimizer(optimizer,
                        model,
                        prefix='',
                        bias_lr_mult=1,
                        bias_decay_mult=1,
                        norm_decay_mult=1,
                        dwconv_decay_mult=1,
                        dcn_offset_lr_mult=1,
                        bypass_duplicate=False):
120
121
122
123
124
125
126
127
128
129
130
    param_groups = optimizer.param_groups
    assert isinstance(optimizer, torch.optim.SGD)
    assert optimizer.defaults['lr'] == base_lr
    assert optimizer.defaults['momentum'] == momentum
    assert optimizer.defaults['weight_decay'] == base_wd
    model_parameters = list(model.parameters())
    assert len(param_groups) == len(model_parameters)
    for i, param in enumerate(model_parameters):
        param_group = param_groups[i]
        assert torch.equal(param_group['params'][0], param)
        assert param_group['momentum'] == momentum
131

132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
    # param1
    param1 = param_groups[0]
    assert param1['lr'] == base_lr
    assert param1['weight_decay'] == base_wd
    # conv1.weight
    conv1_weight = param_groups[1]
    assert conv1_weight['lr'] == base_lr
    assert conv1_weight['weight_decay'] == base_wd
    # conv2.weight
    conv2_weight = param_groups[2]
    assert conv2_weight['lr'] == base_lr
    assert conv2_weight['weight_decay'] == base_wd
    # conv2.bias
    conv2_bias = param_groups[3]
    assert conv2_bias['lr'] == base_lr * bias_lr_mult
    assert conv2_bias['weight_decay'] == base_wd * bias_decay_mult
    # bn.weight
    bn_weight = param_groups[4]
    assert bn_weight['lr'] == base_lr
    assert bn_weight['weight_decay'] == base_wd * norm_decay_mult
    # bn.bias
    bn_bias = param_groups[5]
    assert bn_bias['lr'] == base_lr
    assert bn_bias['weight_decay'] == base_wd * norm_decay_mult
    # sub.param1
    sub_param1 = param_groups[6]
    assert sub_param1['lr'] == base_lr
    assert sub_param1['weight_decay'] == base_wd
    # sub.conv1.weight
    sub_conv1_weight = param_groups[7]
    assert sub_conv1_weight['lr'] == base_lr
    assert sub_conv1_weight['weight_decay'] == base_wd * dwconv_decay_mult
    # sub.conv1.bias
    sub_conv1_bias = param_groups[8]
    assert sub_conv1_bias['lr'] == base_lr * bias_lr_mult
    assert sub_conv1_bias['weight_decay'] == base_wd * dwconv_decay_mult
    # sub.gn.weight
    sub_gn_weight = param_groups[9]
    assert sub_gn_weight['lr'] == base_lr
    assert sub_gn_weight['weight_decay'] == base_wd * norm_decay_mult
    # sub.gn.bias
    sub_gn_bias = param_groups[10]
    assert sub_gn_bias['lr'] == base_lr
    assert sub_gn_bias['weight_decay'] == base_wd * norm_decay_mult

177
178
179
180
181
182
183
184
185
186
187
188
189
    if torch.cuda.is_available():
        dcn_conv_weight = param_groups[11]
        assert dcn_conv_weight['lr'] == base_lr
        assert dcn_conv_weight['weight_decay'] == base_wd

        dcn_offset_weight = param_groups[12]
        assert dcn_offset_weight['lr'] == base_lr * dcn_offset_lr_mult
        assert dcn_offset_weight['weight_decay'] == base_wd

        dcn_offset_bias = param_groups[13]
        assert dcn_offset_bias['lr'] == base_lr * dcn_offset_lr_mult
        assert dcn_offset_bias['weight_decay'] == base_wd

190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270

def test_default_optimizer_constructor():
    model = ExampleModel()

    with pytest.raises(TypeError):
        # optimizer_cfg must be a dict
        optimizer_cfg = []
        optim_constructor = DefaultOptimizerConstructor(optimizer_cfg)
        optim_constructor(model)

    with pytest.raises(TypeError):
        # paramwise_cfg must be a dict or None
        optimizer_cfg = dict(lr=0.0001)
        paramwise_cfg = ['error']
        optim_constructor = DefaultOptimizerConstructor(
            optimizer_cfg, paramwise_cfg)
        optim_constructor(model)

    with pytest.raises(ValueError):
        # bias_decay_mult/norm_decay_mult is specified but weight_decay is None
        optimizer_cfg = dict(lr=0.0001, weight_decay=None)
        paramwise_cfg = dict(bias_decay_mult=1, norm_decay_mult=1)
        optim_constructor = DefaultOptimizerConstructor(
            optimizer_cfg, paramwise_cfg)
        optim_constructor(model)

    # basic config with ExampleModel
    optimizer_cfg = dict(
        type='SGD', lr=base_lr, weight_decay=base_wd, momentum=momentum)
    optim_constructor = DefaultOptimizerConstructor(optimizer_cfg)
    optimizer = optim_constructor(model)
    check_default_optimizer(optimizer, model)

    # basic config with pseudo data parallel
    model = PseudoDataParallel()
    optimizer_cfg = dict(
        type='SGD', lr=base_lr, weight_decay=base_wd, momentum=momentum)
    paramwise_cfg = None
    optim_constructor = DefaultOptimizerConstructor(optimizer_cfg)
    optimizer = optim_constructor(model)
    check_default_optimizer(optimizer, model, prefix='module.')

    # basic config with DataParallel
    if torch.cuda.is_available():
        model = torch.nn.DataParallel(ExampleModel())
        optimizer_cfg = dict(
            type='SGD', lr=base_lr, weight_decay=base_wd, momentum=momentum)
        paramwise_cfg = None
        optim_constructor = DefaultOptimizerConstructor(optimizer_cfg)
        optimizer = optim_constructor(model)
        check_default_optimizer(optimizer, model, prefix='module.')

    # Empty paramwise_cfg with ExampleModel
    model = ExampleModel()
    optimizer_cfg = dict(
        type='SGD', lr=base_lr, weight_decay=base_wd, momentum=momentum)
    paramwise_cfg = dict()
    optim_constructor = DefaultOptimizerConstructor(optimizer_cfg,
                                                    paramwise_cfg)
    optimizer = optim_constructor(model)
    check_default_optimizer(optimizer, model)

    # Empty paramwise_cfg with ExampleModel and no grad
    model = ExampleModel()
    for param in model.parameters():
        param.requires_grad = False
    optimizer_cfg = dict(
        type='SGD', lr=base_lr, weight_decay=base_wd, momentum=momentum)
    paramwise_cfg = dict()
    optim_constructor = DefaultOptimizerConstructor(optimizer_cfg)
    optimizer = optim_constructor(model)
    check_default_optimizer(optimizer, model)

    # paramwise_cfg with ExampleModel
    model = ExampleModel()
    optimizer_cfg = dict(
        type='SGD', lr=base_lr, weight_decay=base_wd, momentum=momentum)
    paramwise_cfg = dict(
        bias_lr_mult=2,
        bias_decay_mult=0.5,
        norm_decay_mult=0,
271
272
        dwconv_decay_mult=0.1,
        dcn_offset_lr_mult=0.1)
273
274
275
    optim_constructor = DefaultOptimizerConstructor(optimizer_cfg,
                                                    paramwise_cfg)
    optimizer = optim_constructor(model)
276
    check_sgd_optimizer(optimizer, model, **paramwise_cfg)
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316

    # paramwise_cfg with ExampleModel, weight decay is None
    model = ExampleModel()
    optimizer_cfg = dict(type='Rprop', lr=base_lr)
    paramwise_cfg = dict(bias_lr_mult=2)
    optim_constructor = DefaultOptimizerConstructor(optimizer_cfg,
                                                    paramwise_cfg)
    optimizer = optim_constructor(model)

    param_groups = optimizer.param_groups
    assert isinstance(optimizer, torch.optim.Rprop)
    assert optimizer.defaults['lr'] == base_lr
    model_parameters = list(model.parameters())
    assert len(param_groups) == len(model_parameters)
    for i, param in enumerate(model_parameters):
        param_group = param_groups[i]
        assert torch.equal(param_group['params'][0], param)
    # param1
    assert param_groups[0]['lr'] == base_lr
    # conv1.weight
    assert param_groups[1]['lr'] == base_lr
    # conv2.weight
    assert param_groups[2]['lr'] == base_lr
    # conv2.bias
    assert param_groups[3]['lr'] == base_lr * paramwise_cfg['bias_lr_mult']
    # bn.weight
    assert param_groups[4]['lr'] == base_lr
    # bn.bias
    assert param_groups[5]['lr'] == base_lr
    # sub.param1
    assert param_groups[6]['lr'] == base_lr
    # sub.conv1.weight
    assert param_groups[7]['lr'] == base_lr
    # sub.conv1.bias
    assert param_groups[8]['lr'] == base_lr * paramwise_cfg['bias_lr_mult']
    # sub.gn.weight
    assert param_groups[9]['lr'] == base_lr
    # sub.gn.bias
    assert param_groups[10]['lr'] == base_lr

317
318
319
320
321
322
323
324
    if OPS_AVAILABLE:
        # dcn.weight
        assert param_groups[11]['lr'] == base_lr
        # dcn.conv_offset.weight
        assert param_groups[12]['lr'] == base_lr
        # dcn.conv_offset.bias
        assert param_groups[13]['lr'] == base_lr

325
326
327
328
329
330
331
332
    # paramwise_cfg with pseudo data parallel
    model = PseudoDataParallel()
    optimizer_cfg = dict(
        type='SGD', lr=base_lr, weight_decay=base_wd, momentum=momentum)
    paramwise_cfg = dict(
        bias_lr_mult=2,
        bias_decay_mult=0.5,
        norm_decay_mult=0,
333
334
        dwconv_decay_mult=0.1,
        dcn_offset_lr_mult=0.1)
335
336
337
    optim_constructor = DefaultOptimizerConstructor(optimizer_cfg,
                                                    paramwise_cfg)
    optimizer = optim_constructor(model)
338
    check_sgd_optimizer(optimizer, model, prefix='module.', **paramwise_cfg)
339
340
341
342
343
344
345
346
347
348

    # paramwise_cfg with DataParallel
    if torch.cuda.is_available():
        model = torch.nn.DataParallel(ExampleModel())
        optimizer_cfg = dict(
            type='SGD', lr=base_lr, weight_decay=base_wd, momentum=momentum)
        paramwise_cfg = dict(
            bias_lr_mult=2,
            bias_decay_mult=0.5,
            norm_decay_mult=0,
349
350
            dwconv_decay_mult=0.1,
            dcn_offset_lr_mult=0.1)
351
352
353
        optim_constructor = DefaultOptimizerConstructor(
            optimizer_cfg, paramwise_cfg)
        optimizer = optim_constructor(model)
354
355
        check_sgd_optimizer(
            optimizer, model, prefix='module.', **paramwise_cfg)
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395

    # paramwise_cfg with ExampleModel and no grad
    for param in model.parameters():
        param.requires_grad = False
    optim_constructor = DefaultOptimizerConstructor(optimizer_cfg,
                                                    paramwise_cfg)
    optimizer = optim_constructor(model)
    param_groups = optimizer.param_groups
    assert isinstance(optimizer, torch.optim.SGD)
    assert optimizer.defaults['lr'] == base_lr
    assert optimizer.defaults['momentum'] == momentum
    assert optimizer.defaults['weight_decay'] == base_wd
    for i, (name, param) in enumerate(model.named_parameters()):
        param_group = param_groups[i]
        assert torch.equal(param_group['params'][0], param)
        assert param_group['momentum'] == momentum
        assert param_group['lr'] == base_lr
        assert param_group['weight_decay'] == base_wd

    # paramwise_cfg with bypass_duplicate option
    model = ExampleDuplicateModel()
    optimizer_cfg = dict(
        type='SGD', lr=base_lr, weight_decay=base_wd, momentum=momentum)
    paramwise_cfg = dict(
        bias_lr_mult=2,
        bias_decay_mult=0.5,
        norm_decay_mult=0,
        dwconv_decay_mult=0.1)
    with pytest.raises(ValueError) as excinfo:
        optim_constructor = DefaultOptimizerConstructor(
            optimizer_cfg, paramwise_cfg)
        optim_constructor(model)
        assert 'some parameters appear in more than one parameter ' \
               'group' == excinfo.value

    paramwise_cfg = dict(
        bias_lr_mult=2,
        bias_decay_mult=0.5,
        norm_decay_mult=0,
        dwconv_decay_mult=0.1,
396
        dcn_offset_lr_mult=0.1,
397
398
399
400
401
402
403
404
405
406
        bypass_duplicate=True)
    optim_constructor = DefaultOptimizerConstructor(optimizer_cfg,
                                                    paramwise_cfg)
    with warnings.catch_warnings(record=True) as w:
        optimizer = optim_constructor(model)
        warnings.simplefilter('always')
        assert len(w) == 1
        assert str(w[0].message) == 'conv3.0 is duplicate. It is skipped ' \
                                    'since bypass_duplicate=True'
    model_parameters = list(model.parameters())
407
408
409
    num_params = 14 if OPS_AVAILABLE else 11
    assert len(optimizer.param_groups) == len(model_parameters) == num_params
    check_sgd_optimizer(optimizer, model, **paramwise_cfg)
410

411
    # test DefaultOptimizerConstructor with custom_keys and ExampleModel
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
    model = ExampleModel()
    optimizer_cfg = dict(
        type='SGD', lr=base_lr, weight_decay=base_wd, momentum=momentum)
    paramwise_cfg = dict(
        custom_keys={
            'param1': dict(lr_mult=10),
            'sub': dict(lr_mult=0.1, decay_mult=0),
            'sub.gn': dict(lr_mult=0.01),
            'non_exist_key': dict(lr_mult=0.0)
        },
        norm_decay_mult=0.5)

    with pytest.raises(TypeError):
        # custom_keys should be a dict
        paramwise_cfg_ = dict(custom_keys=[0.1, 0.0001])
        optim_constructor = DefaultOptimizerConstructor(
428
429
430
431
432
433
434
435
436
437
            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_)
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
        optimizer = optim_constructor(model)

    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'] == base_wd

    # check params groups
    param_groups = optimizer.param_groups

    groups = []
    group_settings = []
    # group 1, matches of 'param1'
    # 'param1' is the longest match for 'sub.param1'
    groups.append(['param1', 'sub.param1'])
    group_settings.append({
        'lr': base_lr * 10,
        'momentum': momentum,
        'weight_decay': base_wd,
    })
    # group 2, matches of 'sub.gn'
    groups.append(['sub.gn.weight', 'sub.gn.bias'])
    group_settings.append({
        'lr': base_lr * 0.01,
        'momentum': momentum,
        'weight_decay': base_wd,
    })
    # group 3, matches of 'sub'
    groups.append(['sub.conv1.weight', 'sub.conv1.bias'])
    group_settings.append({
        'lr': base_lr * 0.1,
        'momentum': momentum,
        'weight_decay': 0,
    })
    # group 4, bn is configured by 'norm_decay_mult'
    groups.append(['bn.weight', 'bn.bias'])
    group_settings.append({
        'lr': base_lr,
        'momentum': momentum,
        'weight_decay': base_wd * 0.5,
    })
    # group 5, default group
    groups.append(['conv1.weight', 'conv2.weight', 'conv2.bias'])
    group_settings.append({
        'lr': base_lr,
        'momentum': momentum,
        'weight_decay': base_wd
    })

491
492
    num_params = 14 if OPS_AVAILABLE else 11
    assert len(param_groups) == num_params
493
494
495
496
497
498
499
500
    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}'

501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
    # 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
    })

538
539
    num_params = 14 if OPS_AVAILABLE else 11
    assert len(param_groups) == num_params
540
541
542
543
544
545
546
547
    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}'

548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564

def test_torch_optimizers():
    torch_optimizers = [
        'ASGD', 'Adadelta', 'Adagrad', 'Adam', 'AdamW', 'Adamax', 'LBFGS',
        'Optimizer', 'RMSprop', 'Rprop', 'SGD', 'SparseAdam'
    ]
    assert set(torch_optimizers).issubset(set(TORCH_OPTIMIZERS))


def test_build_optimizer_constructor():
    model = ExampleModel()
    optimizer_cfg = dict(
        type='SGD', lr=base_lr, weight_decay=base_wd, momentum=momentum)
    paramwise_cfg = dict(
        bias_lr_mult=2,
        bias_decay_mult=0.5,
        norm_decay_mult=0,
565
566
        dwconv_decay_mult=0.1,
        dcn_offset_lr_mult=0.1)
567
568
569
570
571
572
    optim_constructor_cfg = dict(
        type='DefaultOptimizerConstructor',
        optimizer_cfg=optimizer_cfg,
        paramwise_cfg=paramwise_cfg)
    optim_constructor = build_optimizer_constructor(optim_constructor_cfg)
    optimizer = optim_constructor(model)
573
    check_sgd_optimizer(optimizer, model, **paramwise_cfg)
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635

    from mmcv.runner import OPTIMIZERS
    from mmcv.utils import build_from_cfg

    @OPTIMIZER_BUILDERS.register_module()
    class MyOptimizerConstructor(DefaultOptimizerConstructor):

        def __call__(self, model):
            if hasattr(model, 'module'):
                model = model.module

            conv1_lr_mult = self.paramwise_cfg.get('conv1_lr_mult', 1.)

            params = []
            for name, param in model.named_parameters():
                param_group = {'params': [param]}
                if name.startswith('conv1') and param.requires_grad:
                    param_group['lr'] = self.base_lr * conv1_lr_mult
                params.append(param_group)
            optimizer_cfg['params'] = params

            return build_from_cfg(optimizer_cfg, OPTIMIZERS)

    paramwise_cfg = dict(conv1_lr_mult=5)
    optim_constructor_cfg = dict(
        type='MyOptimizerConstructor',
        optimizer_cfg=optimizer_cfg,
        paramwise_cfg=paramwise_cfg)
    optim_constructor = build_optimizer_constructor(optim_constructor_cfg)
    optimizer = optim_constructor(model)

    param_groups = optimizer.param_groups
    assert isinstance(optimizer, torch.optim.SGD)
    assert optimizer.defaults['lr'] == base_lr
    assert optimizer.defaults['momentum'] == momentum
    assert optimizer.defaults['weight_decay'] == base_wd
    for i, param in enumerate(model.parameters()):
        param_group = param_groups[i]
        assert torch.equal(param_group['params'][0], param)
        assert param_group['momentum'] == momentum
    # conv1.weight
    assert param_groups[1]['lr'] == base_lr * paramwise_cfg['conv1_lr_mult']
    assert param_groups[1]['weight_decay'] == base_wd


def test_build_optimizer():
    model = ExampleModel()
    optimizer_cfg = dict(
        type='SGD', lr=base_lr, weight_decay=base_wd, momentum=momentum)
    optimizer = build_optimizer(model, optimizer_cfg)
    check_default_optimizer(optimizer, model)

    model = ExampleModel()
    optimizer_cfg = dict(
        type='SGD',
        lr=base_lr,
        weight_decay=base_wd,
        momentum=momentum,
        paramwise_cfg=dict(
            bias_lr_mult=2,
            bias_decay_mult=0.5,
            norm_decay_mult=0,
636
637
            dwconv_decay_mult=0.1,
            dcn_offset_lr_mult=0.1))
638
    optimizer = build_optimizer(model, optimizer_cfg)
639
    check_sgd_optimizer(optimizer, model, **optimizer_cfg['paramwise_cfg'])