| import torch | |
| import torch.optim as optim | |
| import numpy as np | |
| import itertools | |
| def singleton(class_): | |
| instances = {} | |
| def getinstance(*args, **kwargs): | |
| if class_ not in instances: | |
| instances[class_] = class_(*args, **kwargs) | |
| return instances[class_] | |
| return getinstance | |
| class get_optimizer(object): | |
| def __init__(self): | |
| self.optimizer = {} | |
| self.register(optim.SGD, 'sgd') | |
| self.register(optim.Adam, 'adam') | |
| self.register(optim.AdamW, 'adamw') | |
| def register(self, optim, name): | |
| self.optimizer[name] = optim | |
| def __call__(self, net, cfg): | |
| if cfg is None: | |
| return None | |
| t = cfg.type | |
| if isinstance(net, (torch.nn.DataParallel, | |
| torch.nn.parallel.DistributedDataParallel)): | |
| netm = net.module | |
| else: | |
| netm = net | |
| pg = getattr(netm, 'parameter_group', None) | |
| if pg is not None: | |
| params = [] | |
| for group_name, module_or_para in pg.items(): | |
| if not isinstance(module_or_para, list): | |
| module_or_para = [module_or_para] | |
| grouped_params = [mi.parameters() if isinstance(mi, torch.nn.Module) else [mi] for mi in module_or_para] | |
| grouped_params = itertools.chain(*grouped_params) | |
| pg_dict = {'params': grouped_params, 'name': group_name} | |
| params.append(pg_dict) | |
| else: | |
| params = net.parameters() | |
| return self.optimizer[t](params, lr=0, **cfg.args) | |