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