|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
""" |
|
|
Mostly copy-paste from BEiT library: |
|
|
https://github.com/microsoft/unilm/blob/master/beit/semantic_segmentation/mmcv_custom/layer_decay_optimizer_constructor.py |
|
|
""" |
|
|
|
|
|
import json |
|
|
|
|
|
from mmcv.runner import (OPTIMIZER_BUILDERS, DefaultOptimizerConstructor, |
|
|
get_dist_info) |
|
|
from mmseg.utils import get_root_logger |
|
|
|
|
|
|
|
|
def get_num_layer_for_swin(var_name, num_max_layer, depths): |
|
|
if var_name.startswith('backbone.patch_embed'): |
|
|
return 0 |
|
|
elif var_name.startswith('decode_head.mask_embed'): |
|
|
return 0 |
|
|
elif var_name.startswith('decode_head.cls_embed'): |
|
|
return 0 |
|
|
elif var_name.startswith('decode_head.level_embed'): |
|
|
return 0 |
|
|
elif var_name.startswith('decode_head.query_embed'): |
|
|
return 0 |
|
|
elif var_name.startswith('decode_head.query_feat'): |
|
|
return 0 |
|
|
if var_name.startswith('backbone.cb_modules.0.patch_embed'): |
|
|
return 0 |
|
|
elif 'level_embeds' in var_name: |
|
|
return 0 |
|
|
elif var_name.startswith('backbone.layers') or var_name.startswith( |
|
|
'backbone.levels'): |
|
|
if var_name.split('.')[3] not in ['downsample', 'norm']: |
|
|
stage_id = int(var_name.split('.')[2]) |
|
|
layer_id = int(var_name.split('.')[4]) |
|
|
|
|
|
if stage_id == 0: |
|
|
return layer_id + 1 |
|
|
elif stage_id == 1: |
|
|
return layer_id + 1 + depths[0] |
|
|
elif stage_id == 2: |
|
|
return layer_id + 1 + depths[0] + depths[1] |
|
|
else: |
|
|
return layer_id + 1 + depths[0] + depths[1] + depths[2] |
|
|
else: |
|
|
stage_id = int(var_name.split('.')[2]) |
|
|
if stage_id == 0: |
|
|
return 1 + depths[0] |
|
|
elif stage_id == 1: |
|
|
return 1 + depths[0] + depths[1] |
|
|
elif stage_id == 2: |
|
|
return 1 + depths[0] + depths[1] + depths[2] |
|
|
else: |
|
|
return 1 + depths[0] + depths[1] + depths[2] |
|
|
else: |
|
|
return num_max_layer - 1 |
|
|
|
|
|
|
|
|
@OPTIMIZER_BUILDERS.register_module() |
|
|
class CustomLayerDecayOptimizerConstructor(DefaultOptimizerConstructor): |
|
|
|
|
|
def add_params(self, params, module, prefix='', is_dcn_module=None): |
|
|
"""Add all parameters of module to the params list. |
|
|
The parameters of the given module will be added to the list of param |
|
|
groups, with specific rules defined by paramwise_cfg. |
|
|
Args: |
|
|
params (list[dict]): A list of param groups, it will be modified |
|
|
in place. |
|
|
module (nn.Module): The module to be added. |
|
|
prefix (str): The prefix of the module |
|
|
is_dcn_module (int|float|None): If the current module is a |
|
|
submodule of DCN, `is_dcn_module` will be passed to |
|
|
control conv_offset layer's learning rate. Defaults to None. |
|
|
""" |
|
|
parameter_groups = {} |
|
|
logger = get_root_logger() |
|
|
logger.info(self.paramwise_cfg) |
|
|
backbone_small_lr = self.paramwise_cfg.get('backbone_small_lr', False) |
|
|
dino_head = self.paramwise_cfg.get('dino_head', False) |
|
|
num_layers = self.paramwise_cfg.get('num_layers') + 2 |
|
|
layer_decay_rate = self.paramwise_cfg.get('layer_decay_rate') |
|
|
depths = self.paramwise_cfg.get('depths') |
|
|
offset_lr_scale = self.paramwise_cfg.get('offset_lr_scale', 1.0) |
|
|
|
|
|
logger.info('Build CustomLayerDecayOptimizerConstructor %f - %d' % |
|
|
(layer_decay_rate, num_layers)) |
|
|
weight_decay = self.base_wd |
|
|
|
|
|
for name, param in module.named_parameters(): |
|
|
if not param.requires_grad: |
|
|
continue |
|
|
if len(param.shape) == 1 or name.endswith('.bias') or \ |
|
|
'relative_position' in name or \ |
|
|
'norm' in name or\ |
|
|
'sampling_offsets' in name: |
|
|
group_name = 'no_decay' |
|
|
this_weight_decay = 0. |
|
|
else: |
|
|
group_name = 'decay' |
|
|
this_weight_decay = weight_decay |
|
|
|
|
|
layer_id = get_num_layer_for_swin(name, num_layers, depths) |
|
|
if layer_id == num_layers - 1 and dino_head and \ |
|
|
('sampling_offsets' in name or 'reference_points' in name): |
|
|
group_name = 'layer_%d_%s_0.1x' % (layer_id, group_name) |
|
|
elif ('sampling_offsets' in name or 'reference_points' in name) and 'backbone' in name: |
|
|
group_name = 'layer_%d_%s_offset_lr_scale' % (layer_id, |
|
|
group_name) |
|
|
else: |
|
|
group_name = 'layer_%d_%s' % (layer_id, group_name) |
|
|
|
|
|
if group_name not in parameter_groups: |
|
|
scale = layer_decay_rate ** (num_layers - layer_id - 1) |
|
|
if scale < 1 and backbone_small_lr == True: |
|
|
scale = scale * 0.1 |
|
|
if '0.1x' in group_name: |
|
|
scale = scale * 0.1 |
|
|
if 'offset_lr_scale' in group_name: |
|
|
scale = scale * offset_lr_scale |
|
|
|
|
|
parameter_groups[group_name] = { |
|
|
'weight_decay': this_weight_decay, |
|
|
'params': [], |
|
|
'param_names': [], |
|
|
'lr_scale': scale, |
|
|
'group_name': group_name, |
|
|
'lr': scale * self.base_lr, |
|
|
} |
|
|
|
|
|
parameter_groups[group_name]['params'].append(param) |
|
|
parameter_groups[group_name]['param_names'].append(name) |
|
|
rank, _ = get_dist_info() |
|
|
if rank == 0: |
|
|
to_display = {} |
|
|
for key in parameter_groups: |
|
|
to_display[key] = { |
|
|
'param_names': parameter_groups[key]['param_names'], |
|
|
'lr_scale': parameter_groups[key]['lr_scale'], |
|
|
'lr': parameter_groups[key]['lr'], |
|
|
'weight_decay': parameter_groups[key]['weight_decay'], |
|
|
} |
|
|
logger.info('Param groups = %s' % json.dumps(to_display, indent=2)) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
params.extend(parameter_groups.values()) |
|
|
|
