in easycv/apis/train.py [0:0]
def build_optimizer(model, optimizer_cfg):
"""Build optimizer from configs.
Args:
model (:obj:`nn.Module`): The model with parameters to be optimized.
optimizer_cfg (dict): The config dict of the optimizer.
Positional fields are:
- type: class name of the optimizer.
- lr: base learning rate.
Optional fields are:
- any arguments of the corresponding optimizer type, e.g.,
weight_decay, momentum, etc.
- paramwise_options: a dict with regular expression as keys
to match parameter names and a dict containing options as
values. Options include 6 fields: lr, lr_mult, momentum,
momentum_mult, weight_decay, weight_decay_mult.
Returns:
torch.optim.Optimizer: The initialized optimizer.
Example:
>>> model = torch.nn.modules.Conv1d(1, 1, 1)
>>> paramwise_options = {
>>> '(bn|gn)(\d+)?.(weight|bias)': dict(weight_decay_mult=0.1),
>>> '\Ahead.': dict(lr_mult=10, momentum=0)}
>>> optimizer_cfg = dict(type='SGD', lr=0.01, momentum=0.9,
>>> weight_decay=0.0001,
>>> paramwise_options=paramwise_options)
>>> optimizer = build_optimizer(model, optimizer_cfg)
"""
if hasattr(model, 'module'):
model = model.module
# some special model (DINO) only need to optimize parts of parameter, this kind of model will
# provide attribute get_params_groups to initial optimizer, as we catch this attribute, we do this if
if hasattr(model, 'get_params_groups'):
print('type : ', type(model),
'trigger opimizer model param_groups set for DINO')
parameters = model.get_params_groups()
optimizer_cfg = optimizer_cfg.copy()
optimizer_cls = getattr(optimizer, optimizer_cfg.pop('type'))
return optimizer_cls(parameters, **optimizer_cfg)
# for some model which use transformer(swin/shuffle/cswin), we should set it bias with no weight decay
set_var_bias_nowd = optimizer_cfg.pop('set_var_bias_nowd', None)
if set_var_bias_nowd is None:
set_var_bias_nowd = optimizer_cfg.pop(
'trans_weight_decay_set', None
) # this is failback when we switch version, set_var_bias_nowd used called trans_weight_decay_set
if set_var_bias_nowd is not None:
print('type : ', type(model), 'trigger transformer set_var_bias_nowd')
skip = []
skip_keywords = []
assert (type(set_var_bias_nowd) is list)
for model_part in set_var_bias_nowd:
mpart = getattr(model, model_part, None)
if mpart is not None:
tskip, tskip_keywords = get_skip_list_keywords(mpart)
skip += tskip
skip_keywords += tskip_keywords
parameters = _set_weight_decay(model, skip, skip_keywords)
optimizer_cfg = optimizer_cfg.copy()
optimizer_cls = getattr(optimizer, optimizer_cfg.pop('type'))
return optimizer_cls(parameters, **optimizer_cfg)
constructor_type = optimizer_cfg.pop('constructor', None)
optimizer_cfg = optimizer_cfg.copy()
paramwise_options = optimizer_cfg.pop('paramwise_options', None)
# if no paramwise option is specified, just use the global setting
if constructor_type is not None:
optimizer_cls = getattr(optimizer, optimizer_cfg.pop('type'))
optim_constructor = build_optimizer_constructor(
dict(
type=constructor_type,
optimizer_cfg=optimizer_cfg,
paramwise_cfg=paramwise_options))
params = []
optim_constructor.add_params(params, model)
return optimizer_cls(params, **optimizer_cfg)
elif paramwise_options is None:
return obj_from_dict(optimizer_cfg, optimizer,
dict(params=model.parameters()))
else:
assert isinstance(paramwise_options, dict)
params = []
for name, param in model.named_parameters():
param_group = {'params': [param]}
if not param.requires_grad:
params.append(param_group)
continue
for regexp, options in paramwise_options.items():
if re.search(regexp, name):
for key, value in options.items():
if key.endswith('_mult'): # is a multiplier
key = key[:-5]
assert key in optimizer_cfg, \
'{} not in optimizer_cfg'.format(key)
value = optimizer_cfg[key] * value
param_group[key] = value
if not dist.is_initialized() or dist.get_rank() == 0:
print_log('paramwise_options -- {}: {}={}'.format(
name, key, value))
# otherwise use the global settings
params.append(param_group)
optimizer_cls = getattr(optimizer, optimizer_cfg.pop('type'))
return optimizer_cls(params, **optimizer_cfg)