|
|
| from functools import reduce
|
| from operator import mul
|
| from typing import List
|
|
|
| import torch.nn as nn
|
| from mmengine.logging import print_log
|
| from mmengine.optim import DefaultOptimWrapperConstructor
|
|
|
| from mmaction.registry import OPTIM_WRAPPER_CONSTRUCTORS
|
|
|
|
|
| @OPTIM_WRAPPER_CONSTRUCTORS.register_module()
|
| class SwinOptimWrapperConstructor(DefaultOptimWrapperConstructor):
|
|
|
| def add_params(self,
|
| params: List[dict],
|
| module: nn.Module,
|
| prefix: str = 'base',
|
| **kwargs) -> 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. Defaults to ``'base'``.
|
| """
|
| for name, param in module.named_parameters(recurse=False):
|
| param_group = {'params': [param]}
|
| if not param.requires_grad:
|
| params.append(param_group)
|
| continue
|
|
|
| param_group['lr'] = self.base_lr
|
| if self.base_wd is not None:
|
| param_group['weight_decay'] = self.base_wd
|
|
|
| processing_keys = [
|
| key for key in self.paramwise_cfg if key in f'{prefix}.{name}'
|
| ]
|
| if processing_keys:
|
| param_group['lr'] *= \
|
| reduce(mul, [self.paramwise_cfg[key].get('lr_mult', 1.)
|
| for key in processing_keys])
|
| if self.base_wd is not None:
|
| param_group['weight_decay'] *= \
|
| reduce(mul, [self.paramwise_cfg[key].
|
| get('decay_mult', 1.)
|
| for key in processing_keys])
|
|
|
| params.append(param_group)
|
|
|
| for key, value in param_group.items():
|
| if key == 'params':
|
| continue
|
| full_name = f'{prefix}.{name}' if prefix else name
|
| print_log(
|
| f'paramwise_options -- '
|
| f'{full_name}: {key} = {round(value, 8)}',
|
| logger='current')
|
|
|
| for child_name, child_mod in module.named_children():
|
| child_prefix = f'{prefix}.{child_name}' if prefix else child_name
|
| self.add_params(params, child_mod, prefix=child_prefix)
|
|
|
