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