| |
| import copy |
| import inspect |
|
|
| import torch |
|
|
| from ...utils import Registry, build_from_cfg |
|
|
| OPTIMIZERS = Registry('optimizer') |
| OPTIMIZER_BUILDERS = Registry('optimizer builder') |
|
|
|
|
| def register_torch_optimizers(): |
| torch_optimizers = [] |
| for module_name in dir(torch.optim): |
| if module_name.startswith('__'): |
| continue |
| _optim = getattr(torch.optim, module_name) |
| if inspect.isclass(_optim) and issubclass(_optim, |
| torch.optim.Optimizer): |
| OPTIMIZERS.register_module()(_optim) |
| torch_optimizers.append(module_name) |
| return torch_optimizers |
|
|
|
|
| TORCH_OPTIMIZERS = register_torch_optimizers() |
|
|
|
|
| def build_optimizer_constructor(cfg): |
| return build_from_cfg(cfg, OPTIMIZER_BUILDERS) |
|
|
|
|
| def build_optimizer(model, cfg): |
| optimizer_cfg = copy.deepcopy(cfg) |
| constructor_type = optimizer_cfg.pop('constructor', |
| 'DefaultOptimizerConstructor') |
| paramwise_cfg = optimizer_cfg.pop('paramwise_cfg', None) |
| optim_constructor = build_optimizer_constructor( |
| dict( |
| type=constructor_type, |
| optimizer_cfg=optimizer_cfg, |
| paramwise_cfg=paramwise_cfg)) |
| optimizer = optim_constructor(model) |
| return optimizer |
|
|
