| | |
| | |
| |
|
| |
|
| | import torch |
| | import logging |
| |
|
| | logger = logging.getLogger() |
| |
|
| | class Optimizer(object): |
| | def __init__(self, |
| | model, |
| | lr0, |
| | momentum, |
| | wd, |
| | warmup_steps, |
| | warmup_start_lr, |
| | max_iter, |
| | power, |
| | *args, **kwargs): |
| | self.warmup_steps = warmup_steps |
| | self.warmup_start_lr = warmup_start_lr |
| | self.lr0 = lr0 |
| | self.lr = self.lr0 |
| | self.max_iter = float(max_iter) |
| | self.power = power |
| | self.it = 0 |
| | wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params = model.get_params() |
| | param_list = [ |
| | {'params': wd_params}, |
| | {'params': nowd_params, 'weight_decay': 0}, |
| | {'params': lr_mul_wd_params, 'lr_mul': True}, |
| | {'params': lr_mul_nowd_params, 'weight_decay': 0, 'lr_mul': True}] |
| | self.optim = torch.optim.SGD( |
| | param_list, |
| | lr = lr0, |
| | momentum = momentum, |
| | weight_decay = wd) |
| | self.warmup_factor = (self.lr0/self.warmup_start_lr)**(1./self.warmup_steps) |
| |
|
| |
|
| | def get_lr(self): |
| | if self.it <= self.warmup_steps: |
| | lr = self.warmup_start_lr*(self.warmup_factor**self.it) |
| | else: |
| | factor = (1-(self.it-self.warmup_steps)/(self.max_iter-self.warmup_steps))**self.power |
| | lr = self.lr0 * factor |
| | return lr |
| |
|
| |
|
| | def step(self): |
| | self.lr = self.get_lr() |
| | for pg in self.optim.param_groups: |
| | if pg.get('lr_mul', False): |
| | pg['lr'] = self.lr * 10 |
| | else: |
| | pg['lr'] = self.lr |
| | if self.optim.defaults.get('lr_mul', False): |
| | self.optim.defaults['lr'] = self.lr * 10 |
| | else: |
| | self.optim.defaults['lr'] = self.lr |
| | self.it += 1 |
| | self.optim.step() |
| | if self.it == self.warmup_steps+2: |
| | logger.info('==> warmup done, start to implement poly lr strategy') |
| |
|
| | def zero_grad(self): |
| | self.optim.zero_grad() |
| |
|
| |
|
