Spaces:
Sleeping
Sleeping
# Copyright (c) OpenMMLab. All rights reserved. | |
import sys | |
import warnings | |
from unittest.mock import MagicMock | |
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 | |
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) | |
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() | |
if OPS_AVAILABLE: | |
from mmcv.ops import DeformConv2dPack | |
self.dcn = DeformConv2dPack( | |
3, 4, kernel_size=3, deformable_groups=1) | |
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] | |
if OPS_AVAILABLE: | |
from mmcv.ops import DeformConv2dPack | |
self.dcn = DeformConv2dPack( | |
3, 4, kernel_size=3, deformable_groups=1) | |
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] | |
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' | |
] | |
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]]) | |
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): | |
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 | |
# 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 | |
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 | |
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, | |
dwconv_decay_mult=0.1, | |
dcn_offset_lr_mult=0.1) | |
optim_constructor = DefaultOptimizerConstructor(optimizer_cfg, | |
paramwise_cfg) | |
optimizer = optim_constructor(model) | |
check_sgd_optimizer(optimizer, model, **paramwise_cfg) | |
# 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 | |
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 | |
# 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, | |
dwconv_decay_mult=0.1, | |
dcn_offset_lr_mult=0.1) | |
optim_constructor = DefaultOptimizerConstructor(optimizer_cfg, | |
paramwise_cfg) | |
optimizer = optim_constructor(model) | |
check_sgd_optimizer(optimizer, model, prefix='module.', **paramwise_cfg) | |
# 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, | |
dwconv_decay_mult=0.1, | |
dcn_offset_lr_mult=0.1) | |
optim_constructor = DefaultOptimizerConstructor( | |
optimizer_cfg, paramwise_cfg) | |
optimizer = optim_constructor(model) | |
check_sgd_optimizer( | |
optimizer, model, prefix='module.', **paramwise_cfg) | |
# 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, | |
dcn_offset_lr_mult=0.1, | |
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()) | |
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) | |
# test DefaultOptimizerConstructor with custom_keys and ExampleModel | |
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( | |
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, | |
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 | |
}) | |
num_params = 14 if OPS_AVAILABLE else 11 | |
assert len(param_groups) == num_params | |
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}' | |
# 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 | |
}) | |
num_params = 14 if OPS_AVAILABLE else 11 | |
assert len(param_groups) == num_params | |
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 = [ | |
'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, | |
dwconv_decay_mult=0.1, | |
dcn_offset_lr_mult=0.1) | |
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) | |
check_sgd_optimizer(optimizer, model, **paramwise_cfg) | |
from mmcv.runner import OPTIMIZERS | |
from mmcv.utils import build_from_cfg | |
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, | |
dwconv_decay_mult=0.1, | |
dcn_offset_lr_mult=0.1)) | |
optimizer = build_optimizer(model, optimizer_cfg) | |
check_sgd_optimizer(optimizer, model, **optimizer_cfg['paramwise_cfg']) | |