| | |
| | 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 |
| |
|
