| from functools import partial |
|
|
| import torch |
| from omegaconf import DictConfig |
|
|
|
|
| def get_optimizer(optimizer_name, parameters, lr=None, config: DictConfig = None): |
| param_dict = {} |
| if optimizer_name == "adam": |
| optimizer = partial(torch.optim.Adam, params=parameters) |
| if lr is not None: |
| optimizer = partial(torch.optim.Adam, params=parameters, lr=lr) |
| if config.get('betas'): |
| param_dict['betas'] = config.betas |
| if config.get('weight_decay'): |
| param_dict['weight_decay'] = config.weight_decay |
| if config.get('eps'): |
| param_dict['eps'] = config.eps |
| elif optimizer_name == "adamw": |
| optimizer = partial(torch.optim.AdamW, params=parameters) |
| if lr is not None: |
| optimizer = partial(torch.optim.AdamW, params=parameters, lr=lr) |
| if config.get('betas'): |
| param_dict['betas'] = config.betas |
| if config.get('weight_decay'): |
| param_dict['weight_decay'] = config.weight_decay |
| if config.get('eps'): |
| param_dict['eps'] = config.eps |
| elif optimizer_name == "radam": |
| optimizer = partial(torch.optim.RAdam, params=parameters) |
| if lr is not None: |
| optimizer = partial(torch.optim.RAdam, params=parameters, lr=lr) |
| if config.get('betas'): |
| param_dict['betas'] = config.betas |
| if config.get('weight_decay'): |
| param_dict['weight_decay'] = config.weight_decay |
| elif optimizer_name == "sgd": |
| optimizer = partial(torch.optim.SGD, params=parameters) |
| if lr is not None: |
| optimizer = partial(torch.optim.SGD, params=parameters, lr=lr) |
| if config.get('momentum'): |
| param_dict['momentum'] = config.momentum |
| if config.get('weight_decay'): |
| param_dict['weight_decay'] = config.weight_decay |
| if config.get('nesterov'): |
| param_dict['nesterov'] = config.nesterov |
| else: |
| raise NotImplementedError(f"Optimizer {optimizer_name} not implemented.") |
|
|
| if len(param_dict.keys()) > 0: |
| return optimizer(**param_dict) |
| else: |
| return optimizer() |
|
|
