| |
| import os, sys |
| import os.path as osp |
| import numpy as np |
| import torch |
| from torch import nn |
| from torch.optim import Optimizer |
| from functools import reduce |
| from torch.optim import AdamW |
|
|
| class MultiOptimizer: |
| def __init__(self, optimizers={}, schedulers={}): |
| self.optimizers = optimizers |
| self.schedulers = schedulers |
| self.keys = list(optimizers.keys()) |
| self.param_groups = reduce(lambda x,y: x+y, [v.param_groups for v in self.optimizers.values()]) |
|
|
| def state_dict(self): |
| state_dicts = [(key, self.optimizers[key].state_dict())\ |
| for key in self.keys] |
| return state_dicts |
|
|
| def scheduler_state_dict(self): |
| state_dicts = [(key, self.schedulers[key].state_dict())\ |
| for key in self.keys] |
| return state_dicts |
|
|
| def load_state_dict(self, state_dict): |
| for key, val in state_dict: |
| try: |
| self.optimizers[key].load_state_dict(val) |
| except: |
| print("Unloaded %s" % key) |
|
|
| def load_scheduler_state_dict(self, state_dict): |
| for key, val in state_dict: |
| try: |
| self.schedulers[key].load_state_dict(val) |
| except: |
| print("Unloaded %s" % key) |
|
|
| def step(self, key=None, scaler=None): |
| keys = [key] if key is not None else self.keys |
| _ = [self._step(key, scaler) for key in keys] |
|
|
| def _step(self, key, scaler=None): |
| if scaler is not None: |
| scaler.step(self.optimizers[key]) |
| scaler.update() |
| else: |
| self.optimizers[key].step() |
|
|
| def zero_grad(self, key=None): |
| if key is not None: |
| self.optimizers[key].zero_grad() |
| else: |
| _ = [self.optimizers[key].zero_grad() for key in self.keys] |
|
|
| def scheduler(self, *args, key=None): |
| if key is not None: |
| self.schedulers[key].step(*args) |
| else: |
| _ = [self.schedulers[key].step_batch(*args) for key in self.keys] |
|
|
| def define_scheduler(optimizer, params): |
| scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=params['gamma']) |
|
|
| return scheduler |
|
|
| def build_optimizer(model_dict, lr, type='AdamW'): |
| optim = {} |
| for key, model in model_dict.items(): |
| model_parameters = model.parameters() |
| parameters_names = [] |
| parameters_names.append( |
| [ |
| name_param_pair[0] |
| for name_param_pair in model.named_parameters() |
| ] |
| ) |
| if type == 'AdamW': |
| optim[key] = AdamW( |
| model_parameters, |
| lr=lr, |
| betas=(0.9, 0.98), |
| eps=1e-9, |
| weight_decay=0.1, |
| ) |
| else: |
| raise ValueError('Unknown optimizer type: %s' % type) |
|
|
| schedulers = dict([(key, torch.optim.lr_scheduler.ExponentialLR(opt, gamma=0.999996)) |
| for key, opt in optim.items()]) |
|
|
| multi_optim = MultiOptimizer(optim, schedulers) |
| return multi_optim |
