|
|
|
|
|
import torch
|
|
|
from mmengine.optim import DefaultOptimWrapperConstructor
|
|
|
from mmengine.utils.dl_utils.parrots_wrapper import (SyncBatchNorm_,
|
|
|
_BatchNorm, _ConvNd)
|
|
|
|
|
|
from mmaction.registry import OPTIM_WRAPPER_CONSTRUCTORS
|
|
|
|
|
|
|
|
|
@OPTIM_WRAPPER_CONSTRUCTORS.register_module()
|
|
|
class TSMOptimWrapperConstructor(DefaultOptimWrapperConstructor):
|
|
|
"""Optimizer constructor in TSM model.
|
|
|
|
|
|
This constructor builds optimizer in different ways from the default one.
|
|
|
|
|
|
1. Parameters of the first conv layer have default lr and weight decay.
|
|
|
2. Parameters of BN layers have default lr and zero weight decay.
|
|
|
3. If the field "fc_lr5" in paramwise_cfg is set to True, the parameters
|
|
|
of the last fc layer in cls_head have 5x lr multiplier and 10x weight
|
|
|
decay multiplier.
|
|
|
4. Weights of other layers have default lr and weight decay, and biases
|
|
|
have a 2x lr multiplier and zero weight decay.
|
|
|
"""
|
|
|
|
|
|
def add_params(self, params, model, **kwargs):
|
|
|
"""Add parameters and their corresponding lr and wd to the params.
|
|
|
|
|
|
Args:
|
|
|
params (list): The list to be modified, containing all parameter
|
|
|
groups and their corresponding lr and wd configurations.
|
|
|
model (nn.Module): The model to be trained with the optimizer.
|
|
|
"""
|
|
|
|
|
|
|
|
|
fc_lr5 = self.paramwise_cfg['fc_lr5']
|
|
|
first_conv_weight = []
|
|
|
first_conv_bias = []
|
|
|
normal_weight = []
|
|
|
normal_bias = []
|
|
|
lr5_weight = []
|
|
|
lr10_bias = []
|
|
|
bn = []
|
|
|
|
|
|
conv_cnt = 0
|
|
|
|
|
|
for m in model.modules():
|
|
|
if isinstance(m, _ConvNd):
|
|
|
m_params = list(m.parameters())
|
|
|
conv_cnt += 1
|
|
|
if conv_cnt == 1:
|
|
|
first_conv_weight.append(m_params[0])
|
|
|
if len(m_params) == 2:
|
|
|
first_conv_bias.append(m_params[1])
|
|
|
else:
|
|
|
normal_weight.append(m_params[0])
|
|
|
if len(m_params) == 2:
|
|
|
normal_bias.append(m_params[1])
|
|
|
elif isinstance(m, torch.nn.Linear):
|
|
|
m_params = list(m.parameters())
|
|
|
normal_weight.append(m_params[0])
|
|
|
if len(m_params) == 2:
|
|
|
normal_bias.append(m_params[1])
|
|
|
elif isinstance(m,
|
|
|
(_BatchNorm, SyncBatchNorm_, torch.nn.GroupNorm)):
|
|
|
for param in list(m.parameters()):
|
|
|
if param.requires_grad:
|
|
|
bn.append(param)
|
|
|
elif len(m._modules) == 0:
|
|
|
if len(list(m.parameters())) > 0:
|
|
|
raise ValueError(f'New atomic module type: {type(m)}. '
|
|
|
'Need to give it a learning policy')
|
|
|
|
|
|
|
|
|
last_fc_weight = normal_weight.pop()
|
|
|
last_fc_bias = normal_bias.pop()
|
|
|
if fc_lr5:
|
|
|
lr5_weight.append(last_fc_weight)
|
|
|
lr10_bias.append(last_fc_bias)
|
|
|
else:
|
|
|
normal_weight.append(last_fc_weight)
|
|
|
normal_bias.append(last_fc_bias)
|
|
|
|
|
|
params.append({
|
|
|
'params': first_conv_weight,
|
|
|
'lr': self.base_lr,
|
|
|
'weight_decay': self.base_wd
|
|
|
})
|
|
|
params.append({
|
|
|
'params': first_conv_bias,
|
|
|
'lr': self.base_lr * 2,
|
|
|
'weight_decay': 0
|
|
|
})
|
|
|
params.append({
|
|
|
'params': normal_weight,
|
|
|
'lr': self.base_lr,
|
|
|
'weight_decay': self.base_wd
|
|
|
})
|
|
|
params.append({
|
|
|
'params': normal_bias,
|
|
|
'lr': self.base_lr * 2,
|
|
|
'weight_decay': 0
|
|
|
})
|
|
|
params.append({'params': bn, 'lr': self.base_lr, 'weight_decay': 0})
|
|
|
params.append({
|
|
|
'params': lr5_weight,
|
|
|
'lr': self.base_lr * 5,
|
|
|
'weight_decay': self.base_wd
|
|
|
})
|
|
|
params.append({
|
|
|
'params': lr10_bias,
|
|
|
'lr': self.base_lr * 10,
|
|
|
'weight_decay': 0
|
|
|
})
|
|
|
|
