BIT2024
1 year ago
1 changed files with 729 additions and 0 deletions
Split View
Diff Options
-
+729 -0optim.py
+ 729
- 0
optim.py
View File
| @@ -0,0 +1,729 @@ | |||
| # *************************************************************** | |||
| # Copyright (c) 2023 Jittor. All Rights Reserved. | |||
| # Maintainers: | |||
| # Guowei Yang <471184555@qq.com> | |||
| # Guoye Yang <498731903@qq.com> | |||
| # Wenyang Zhou <576825820@qq.com> | |||
| # Meng-Hao Guo <guomenghao1997@gmail.com> | |||
| # Dun Liang <randonlang@gmail.com>. | |||
| # | |||
| # | |||
| # This file is subject to the terms and conditions defined in | |||
| # file 'LICENSE.txt', which is part of this source code package. | |||
| # *************************************************************** | |||
| import jittor as jt | |||
| import numpy as np | |||
| class Optimizer(object): | |||
| """ Basic class of Optimizer. | |||
| Example:: | |||
| optimizer = nn.SGD(model.parameters(), lr) | |||
| optimizer.step(loss) | |||
| """ | |||
| def __init__(self, params, lr, param_sync_iter=10000): | |||
| self.param_groups = [] | |||
| self.lr = lr | |||
| self.param_sync_iter = param_sync_iter | |||
| assert len(params) > 0, "Length of parameters should not be zero" | |||
| if not isinstance(params[0], dict): | |||
| params = [{'params': params}] | |||
| for pg in params: | |||
| assert isinstance(pg, dict) | |||
| self.param_groups.append(pg) | |||
| self.n_step = 0 | |||
| # __zero_grad is a value for fast determ the grad is zero or not | |||
| # so we can omit 0+x | |||
| self.__zero_grad = True | |||
| self._grad_map = {} | |||
| def add_param_group(self, group): | |||
| self.param_groups.append(group) | |||
| def clip_grad_norm(self, max_norm:float, norm_type:int=2): | |||
| r"""Clips gradient norm of this optimizer. | |||
| The norm is computed over all gradients together. | |||
| Args: | |||
| max_norm (float or int): max norm of the gradients | |||
| norm_type (int): 1-norm or 2-norm | |||
| Example:: | |||
| a = jt.ones(2) | |||
| opt = jt.optim.SGD([a], 0.1) | |||
| loss = a*a | |||
| opt.zero_grad() | |||
| opt.backward(loss) | |||
| print(opt.param_groups[0]['grads'][0].norm()) # output: 2.83 | |||
| opt.clip_grad_norm(0.01, 2) | |||
| print(opt.param_groups[0]['grads'][0].norm()) # output: 0.01 | |||
| opt.step() | |||
| """ | |||
| if self.__zero_grad: return | |||
| grads = [] | |||
| for pg in self.param_groups: | |||
| for p, g in zip(pg["params"], pg["grads"]): | |||
| if p.is_stop_grad(): continue | |||
| grads.append(g.flatten()) | |||
| if len(grads) == 0: return | |||
| total_norm = jt.norm(jt.concat(grads), norm_type) | |||
| clip_coef = jt.minimum(max_norm / (total_norm + 1e-6), 1.0) | |||
| for pg in self.param_groups: | |||
| for p, g in zip(pg["params"], pg["grads"]): | |||
| if p.is_stop_grad(): continue | |||
| g.update(g*clip_coef) | |||
| @property | |||
| def defaults(self): | |||
| exclude = set(("defaults", "pre_step", "step")) | |||
| return { k:v for k, v in self.__dict__.items() | |||
| if k[0] != '_' and k not in exclude and not callable(v) } | |||
| def state_dict(self): | |||
| state = {"defaults": self.defaults} | |||
| return state | |||
| def load_state_dict(self, state): | |||
| def dfs(x): | |||
| if isinstance(x, list): | |||
| for i in range(len(x)): | |||
| x[i] = dfs(x[i]) | |||
| elif isinstance(x, dict): | |||
| for k in x: | |||
| x[k] = dfs(x[k]) | |||
| elif isinstance(x, np.ndarray): | |||
| return jt.array(x).stop_grad() | |||
| elif isinstance(x, jt.Var): | |||
| return x.stop_grad() | |||
| return x | |||
| exclude = set(("param_groups", "params")) | |||
| for k, v in state["defaults"].items(): | |||
| if k not in exclude: | |||
| setattr(self, k, dfs(v)) | |||
| param_groups = dfs(state["defaults"].get('param_groups', None)) | |||
| if param_groups is not None: | |||
| exclude = set(("params",)) | |||
| for i in range(len(param_groups)): | |||
| for k, v in param_groups[i].items(): | |||
| if k not in exclude: | |||
| self.param_groups[i][k] = v | |||
| def zero_grad(self): | |||
| self.__zero_grad = True | |||
| def backward(self, loss, retain_graph=False): | |||
| ''' | |||
| optimize.backward(loss) is used for accumulate multiple step, | |||
| it can be used as following: | |||
| Origin source code :: | |||
| n_iter = 10000 | |||
| batch_size = 100 | |||
| ... | |||
| for i in range(n_iter): | |||
| ... | |||
| loss = calc_loss() | |||
| optimizer.step(loss) | |||
| Accumulation version :: | |||
| n_iter = 10000 | |||
| batch_size = 100 | |||
| accumulation_steps = 10 | |||
| n_iter *= accumulation_steps | |||
| batch_size //= accumulation_steps | |||
| ... | |||
| for i in range(n_iter): | |||
| ... | |||
| loss = calc_loss() | |||
| # if loss is a mean across batch, we need to divide accumulation_steps | |||
| optimizer.backward(loss / accumulation_steps) | |||
| if (i+1) % accumulation_steps == 0: | |||
| optimizer.step() | |||
| ''' | |||
| # clean prev grads | |||
| params = [] | |||
| params_has_grad = [] | |||
| for pg in self.param_groups: | |||
| for p in pg['params']: | |||
| params.append(p) | |||
| if not p.is_stop_grad(): | |||
| params_has_grad.append(p) | |||
| # sync prev params | |||
| jt.sync(params_has_grad) | |||
| # get gradient | |||
| grads = jt.grad(loss, params_has_grad, retain_graph) | |||
| # sync grads and model if in mpi | |||
| if jt.in_mpi: | |||
| dep = [] | |||
| def add_dep(v): | |||
| nonlocal dep | |||
| v._add_dependency(dep) | |||
| dep = [v] | |||
| for g in grads: | |||
| g.assign(g.mpi_all_reduce("mean")) | |||
| add_dep(g._input(0)) | |||
| if self.n_step % self.param_sync_iter == 0: | |||
| for p in params: | |||
| p.assign(p.mpi_broadcast()) | |||
| add_dep(p) | |||
| self.n_step += 1 | |||
| # set up grads in param_groups | |||
| pid = 0 | |||
| for pg in self.param_groups: | |||
| if "grads" not in pg: | |||
| pg["grads"] = [ jt.zeros_like(p).stop_grad().stop_fuse() for p in pg['params'] ] | |||
| pg_grads = pg["grads"] | |||
| for i, p in enumerate(pg['params']): | |||
| if not p.is_stop_grad(): | |||
| # accumulate grad and stop grad of grad | |||
| g = grads[pid].stop_grad() | |||
| if not self.__zero_grad: | |||
| g = g + pg_grads[i] | |||
| pg_grads[i].update(g) | |||
| pid += 1 | |||
| self.__zero_grad = False | |||
| def pre_step(self, loss, retain_graph=False): | |||
| """ something should be done before step, such as calc gradients, mpi sync, and so on. | |||
| Example:: | |||
| class MyOptimizer(Optimizer): | |||
| def step(self, loss): | |||
| self.pre_step(loss) | |||
| ... | |||
| self.post_step() | |||
| """ | |||
| if loss is not None: | |||
| self.backward(loss, retain_graph) | |||
| jt.flags.node_order = 1 | |||
| def post_step(self): | |||
| """ something should be done before step, such as zero grad, and so on. | |||
| Example:: | |||
| class MyOptimizer(Optimizer): | |||
| def step(self, loss): | |||
| self.pre_step(loss) | |||
| ... | |||
| self.post_step() | |||
| """ | |||
| jt.flags.node_order = 0 | |||
| self.zero_grad() | |||
| def step(self, loss=None, retain_graph=False): | |||
| self.pre_step(loss, retain_graph) | |||
| for pg in self.param_groups: | |||
| lr = pg.get("lr", self.lr) | |||
| for p, g in zip(pg["params"], pg["grads"]): | |||
| if p.is_stop_grad(): continue | |||
| p.update(p - g * lr) | |||
| self.post_step() | |||
| def _build_grad_map(self): | |||
| _grad_map = {} | |||
| for pg in self.param_groups: | |||
| for p, g in zip(pg["params"], pg["grads"]): | |||
| _grad_map[id(p)] = g | |||
| self._grad_map = _grad_map | |||
| def find_grad(self, v:jt.Var) -> jt.Var: | |||
| if id(v) not in self._grad_map: | |||
| self._build_grad_map() | |||
| if id(v) not in self._grad_map: | |||
| raise RuntimeError("This variable is not managed by this optimizer") | |||
| return self._grad_map[id(v)] | |||
| def opt_grad(v:jt.Var, opt:Optimizer): | |||
| ''' Get grad of certain variable in optimizer, Example:: | |||
| model = Model() | |||
| optimizer = SGD(model.parameters()) | |||
| ... | |||
| optimizer.backward(loss) | |||
| for p in model.parameters(): | |||
| grad = p.opt_grad(optimizer) | |||
| ''' | |||
| return opt.find_grad(v) | |||
| jt.Var.opt_grad = opt_grad | |||
| class SGD(Optimizer): | |||
| """ SGD Optimizer. | |||
| Example:: | |||
| optimizer = nn.SGD(model.parameters(), lr, momentum=0.9) | |||
| optimizer.step(loss) | |||
| """ | |||
| def __init__(self, params, lr, momentum=0, weight_decay=0, dampening=0, nesterov=False): | |||
| super().__init__(params, lr) | |||
| self.momentum = momentum | |||
| self.weight_decay = weight_decay | |||
| self.dampening = dampening | |||
| self.nesterov = nesterov | |||
| # initialize required arguments | |||
| for pg in self.param_groups: | |||
| values = pg["values"] = [] | |||
| for p in pg["params"]: | |||
| values.append(jt.zeros(p.shape, p.dtype).stop_grad()) | |||
| def add_param_group(self, group): | |||
| values = group["values"] = [] | |||
| for p in group["params"]: | |||
| values.append(jt.zeros(p.shape, p.dtype).stop_grad()) | |||
| self.param_groups.append(group) | |||
| def step(self, loss=None, retain_graph=False): | |||
| self.pre_step(loss, retain_graph=False) | |||
| jt.flags.node_order = 1 | |||
| for pg in self.param_groups: | |||
| # get arguments from each param_groups | |||
| lr = pg.get("lr", self.lr) | |||
| momentum = pg.get("momentum", self.momentum) | |||
| weight_decay = pg.get("weight_decay", self.weight_decay) | |||
| dampening = pg.get("dampening", self.dampening) | |||
| nesterov = pg.get("nesterov", self.nesterov) | |||
| # optimize main body | |||
| for p, g, v in zip(pg["params"], pg["grads"], pg["values"]): | |||
| if p.is_stop_grad(): continue | |||
| dp = p * weight_decay + g | |||
| v.update(momentum * v + dp * (1 - dampening)) | |||
| if nesterov: | |||
| p.update(p - (dp + momentum * v) * lr) | |||
| else: | |||
| p.update(p - v * lr) | |||
| self.post_step() | |||
| class RMSprop(Optimizer): | |||
| """ RMSprop Optimizer. | |||
| Args: | |||
| params(list): parameters of model. | |||
| lr(float): learning rate. | |||
| eps(float): term added to the denominator to avoid division by zero, default 1e-8. | |||
| alpha(float): smoothing constant, default 0.99. | |||
| Example: | |||
| optimizer = nn.RMSprop(model.parameters(), lr) | |||
| optimizer.step(loss) | |||
| """ | |||
| def __init__(self, params, lr=1e-2, eps=1e-8, alpha=0.99): | |||
| super().__init__(params, lr) | |||
| self.eps = eps | |||
| self.alpha = alpha | |||
| # initialize required arguments for each param_groups | |||
| for pg in self.param_groups: | |||
| values = pg["values"] = [] | |||
| for p in pg["params"]: | |||
| values.append(jt.zeros(p.shape, p.dtype).stop_grad()) | |||
| def add_param_group(self, group): | |||
| values = group["values"] = [] | |||
| for p in group["params"]: | |||
| values.append(jt.zeros(p.shape, p.dtype).stop_grad()) | |||
| self.param_groups.append(group) | |||
| def step(self, loss=None, retain_graph=False): | |||
| self.pre_step(loss, retain_graph) | |||
| for pg in self.param_groups: | |||
| # get arguments from each param_groups | |||
| lr = pg.get("lr", self.lr) | |||
| eps = pg.get("eps", self.eps) | |||
| alpha = pg.get("alpha", self.alpha) | |||
| for p, g, v in zip(pg["params"], pg["grads"], pg["values"]): | |||
| if p.is_stop_grad(): continue | |||
| v.update(alpha * v + (1-alpha) * g * g) | |||
| p.update(p - lr * g / (jt.sqrt(v) + eps)) | |||
| self.post_step() | |||
| class Adam(Optimizer): | |||
| """ Adam Optimizer. | |||
| Example:: | |||
| optimizer = nn.Adam(model.parameters(), lr, eps=1e-8, betas=(0.9, 0.999)) | |||
| optimizer.step(loss) | |||
| """ | |||
| def __init__(self, params, lr, eps=1e-8, betas=(0.9, 0.999), weight_decay=0): | |||
| super().__init__(params, lr) | |||
| self.eps = eps | |||
| self.betas = betas | |||
| self.weight_decay = weight_decay | |||
| # assert weight_decay==0, "weight_decay is not supported yet" | |||
| # initialize required arguments for each param_groups | |||
| for pg in self.param_groups: | |||
| values = pg["values"] = [] | |||
| m = pg["m"] = [] | |||
| for p in pg["params"]: | |||
| values.append(jt.zeros(p.shape, p.dtype).stop_grad()) | |||
| m.append(jt.zeros(p.shape, p.dtype).stop_grad()) | |||
| def add_param_group(self, group): | |||
| values = group["values"] = [] | |||
| m = group["m"] = [] | |||
| for p in group["params"]: | |||
| values.append(jt.zeros(p.shape, p.dtype).stop_grad()) | |||
| m.append(jt.zeros(p.shape, p.dtype).stop_grad()) | |||
| self.param_groups.append(group) | |||
| def step(self, loss=None, retain_graph=False): | |||
| self.pre_step(loss, retain_graph) | |||
| n = float(self.n_step) | |||
| jt.flags.node_order = 1 | |||
| for pg in self.param_groups: | |||
| # get arguments from each param_groups | |||
| lr = pg.get("lr", self.lr) | |||
| eps = pg.get("eps", self.eps) | |||
| weight_decay = pg.get("weight_decay", self.weight_decay) | |||
| b0, b1 = pg.get("betas", self.betas) | |||
| for p, g, v, m in zip(pg["params"], pg["grads"], pg["values"], pg["m"]): | |||
| if p.is_stop_grad(): continue | |||
| g = p * weight_decay + g | |||
| m.update(b0 * m + (1-b0) * g) | |||
| v.update(b1 * v + (1-b1) * g * g) | |||
| step_size = lr * jt.sqrt(1-b1**n) / (1-b0 ** n) | |||
| p.update(p - m * step_size / (jt.sqrt(v) + eps)) | |||
| self.post_step() | |||
| class AdamW(Optimizer): | |||
| """ AdamW Optimizer. | |||
| Example:: | |||
| optimizer = nn.AdamW(model.parameters(), lr, eps=1e-8, betas=(0.9, 0.999)) | |||
| optimizer.step(loss) | |||
| """ | |||
| def __init__(self, params, lr, eps=1e-8, betas=(0.9, 0.999), weight_decay=0): | |||
| super().__init__(params, lr) | |||
| self.eps = eps | |||
| self.betas = betas | |||
| self.weight_decay = weight_decay | |||
| # assert weight_decay==0, "weight_decay is not supported yet" | |||
| # initialize required arguments for each param_groups | |||
| for pg in self.param_groups: | |||
| values = pg["values"] = [] | |||
| m = pg["m"] = [] | |||
| for p in pg["params"]: | |||
| values.append(jt.zeros(p.shape, p.dtype).stop_grad()) | |||
| m.append(jt.zeros(p.shape, p.dtype).stop_grad()) | |||
| def add_param_group(self, group): | |||
| values = group["values"] = [] | |||
| m = group["m"] = [] | |||
| for p in group["params"]: | |||
| values.append(jt.zeros(p.shape, p.dtype).stop_grad()) | |||
| m.append(jt.zeros(p.shape, p.dtype).stop_grad()) | |||
| self.param_groups.append(group) | |||
| def step(self, loss=None, retain_graph=False): | |||
| self.pre_step(loss, retain_graph) | |||
| n = float(self.n_step) | |||
| for pg in self.param_groups: | |||
| # get arguments from each param_groups | |||
| lr = pg.get("lr", self.lr) | |||
| eps = pg.get("eps", self.eps) | |||
| weight_decay = pg.get("weight_decay", self.weight_decay) | |||
| b0, b1 = pg.get("betas", self.betas) | |||
| for p, g, v, m in zip(pg["params"], pg["grads"], pg["values"], pg["m"]): | |||
| if p.is_stop_grad(): continue | |||
| p.update(p * (1 - lr * weight_decay)) | |||
| bias_correction1 = 1 - b0 ** n | |||
| bias_correction2 = 1 - b1 ** n | |||
| m.update(b0 * m + (1-b0) * g) #exp_avg | |||
| v.update(b1 * v + (1-b1) * g * g) #exp_avg_sq | |||
| denom = jt.sqrt(v) / jt.sqrt(bias_correction2) + eps | |||
| step_size = lr / bias_correction1 | |||
| p.update(p - step_size * m / denom) | |||
| self.post_step() | |||
| class Adan(Optimizer): | |||
| """ Adan Optimizer. | |||
| Adan was proposed in | |||
| Adan: Adaptive Nesterov Momentum Algorithm for | |||
| Faster Optimizing Deep Models[J].arXiv preprint arXiv:2208.06677, 2022. | |||
| https://arxiv.org/abs/2208.06677 | |||
| Adan is an efficient optimizer for most DNN frameworks: | |||
| - About 2x fewer computational load than SOTAs | |||
| - Robust to training setting and batch size | |||
| - Easy to Plug-and-play | |||
| Arguments: | |||
| params (iterable): iterable of parameters to optimize or | |||
| dicts defining parameter groups. | |||
| lr (float, optional): learning rate. (default: 1e-3) | |||
| betas (Tuple[float, float, flot], optional): coefficients used for | |||
| first- and second-order moments. (default: (0.98, 0.92, 0.99)) | |||
| eps (float, optional): term added to the denominator to improve | |||
| numerical stability. (default: 1e-8) | |||
| weight_decay (float, optional): decoupled weight decay | |||
| (L2 penalty) (default: 0) | |||
| max_grad_norm (float, optional): value used to clip | |||
| global grad norm (default: 0.0 no clip) | |||
| """ | |||
| def __init__(self, params, lr=1e-3, betas=(0.98, 0.92, 0.99), | |||
| eps=1e-8, weight_decay=0.0, max_grad_norm=0.0): | |||
| super().__init__(params, lr) | |||
| self.betas = betas | |||
| self.eps = eps | |||
| self.weight_decay = weight_decay | |||
| self.max_grad_norm = max_grad_norm | |||
| for pg in self.param_groups: | |||
| pg["m"] = [] | |||
| pg["v"] = [] | |||
| pg["d"] = [] | |||
| pg["pre_grad"] = [] | |||
| for p in pg["params"]: | |||
| pg["m"].append(jt.zeros(p.shape, p.dtype).stop_grad()) | |||
| pg["v"].append(jt.zeros(p.shape, p.dtype).stop_grad()) | |||
| pg["d"].append(jt.zeros(p.shape, p.dtype).stop_grad()) | |||
| pg["pre_grad"].append(jt.zeros(p.shape, p.dtype).stop_grad()) | |||
| def add_param_group(self, group): | |||
| group["m"] = [] | |||
| group["v"] = [] | |||
| group["d"] = [] | |||
| group["pre_grad"] = [] | |||
| for p in group["params"]: | |||
| group["m"].append(jt.zeros(p.shape, p.dtype).stop_grad()) | |||
| group["v"].append(jt.zeros(p.shape, p.dtype).stop_grad()) | |||
| group["d"].append(jt.zeros(p.shape, p.dtype).stop_grad()) | |||
| group["pre_grad"].append(jt.zeros(p.shape, p.dtype).stop_grad()) | |||
| self.param_groups.append(group) | |||
| def step(self, loss=None, retain_graph=False): | |||
| self.pre_step(loss, retain_graph) | |||
| n = float(self.n_step) | |||
| for pg in self.param_groups: | |||
| lr = pg.get("lr", self.lr) | |||
| betas = pg.get("betas", self.betas) | |||
| eps = pg.get("eps", self.eps) | |||
| weight_decay = pg.get("weight_decay", self.weight_decay) | |||
| max_grad_norm = pg.get("max_grad_norm", self.max_grad_norm) | |||
| if max_grad_norm>0: self.clip_grad_norm(max_grad_norm) | |||
| beta1, beta2, beta3 = betas | |||
| bias_correction1 = 1 - beta1 ** n | |||
| bias_correction2 = 1 - beta2 ** n | |||
| bias_correction3_sqrt = jt.sqrt(1 - beta3 ** n) | |||
| step_size_diff = lr * beta2 * bias_correction3_sqrt / bias_correction2 | |||
| step_size = lr * bias_correction3_sqrt / bias_correction1 | |||
| eps_bias_sqrt = eps * bias_correction3_sqrt | |||
| for p, g, m, v, d, pre_g in zip(pg["params"], | |||
| pg["grads"], | |||
| pg["m"], | |||
| pg["v"], | |||
| pg["d"], | |||
| pg["pre_grad"]): | |||
| if p.is_stop_grad(): continue | |||
| if self.n_step>0: | |||
| pre_g.update(g - pre_g) # Update pre_g as grad_diff | |||
| m.update(beta1 * m + (1 - beta1) * g) | |||
| d.update(beta2 * d + (1 - beta2) * pre_g) # Use pre_g as grad_diff | |||
| pre_g.update(jt.multiply(pre_g, beta2) + g) # Update pre_g as update (g + beta2 * grad_diff) | |||
| v.update(beta3 * v + (1 - beta3) * pre_g * pre_g) # Use pre_g as update | |||
| p.update(p - (step_size * m + step_size_diff * d) / (jt.sqrt(v) + eps_bias_sqrt)) | |||
| p.update(p / (1 + lr * weight_decay)) | |||
| pre_g.update(g) # Update pre_g for the next iteration | |||
| self.post_step() | |||
| class AdanBelief(Optimizer): | |||
| """ Adan Optimizer. | |||
| Adan was proposed in | |||
| Adan: Adaptive Nesterov Momentum Algorithm for | |||
| Faster Optimizing Deep Models[J].arXiv preprint arXiv:2208.06677, 2022. | |||
| https://arxiv.org/abs/2208.06677 | |||
| Adan is an efficient optimizer for most DNN frameworks: | |||
| - About 2x fewer computational load than SOTAs | |||
| - Robust to training setting and batch size | |||
| - Easy to Plug-and-play | |||
| Arguments: | |||
| params (iterable): iterable of parameters to optimize or | |||
| dicts defining parameter groups. | |||
| lr (float, optional): learning rate. (default: 1e-3) | |||
| betas (Tuple[float, float, flot], optional): coefficients used for | |||
| first- and second-order moments. (default: (0.98, 0.92, 0.99)) | |||
| eps (float, optional): term added to the denominator to improve | |||
| numerical stability. (default: 1e-8) | |||
| weight_decay (float, optional): decoupled weight decay | |||
| (L2 penalty) (default: 0) | |||
| max_grad_norm (float, optional): value used to clip | |||
| global grad norm (default: 0.0 no clip) | |||
| """ | |||
| def __init__(self, params, lr=1e-3, betas=(0.98, 0.92, 0.99), | |||
| eps=1e-8, weight_decay=0.0, max_grad_norm=0.0): | |||
| super().__init__(params, lr) | |||
| self.betas = betas | |||
| self.eps = eps | |||
| self.weight_decay = weight_decay | |||
| self.max_grad_norm = max_grad_norm | |||
| for pg in self.param_groups: | |||
| pg["m"] = [] | |||
| pg["v"] = [] | |||
| pg["d"] = [] | |||
| pg["pre_grad"] = [] | |||
| for p in pg["params"]: | |||
| pg["m"].append(jt.zeros(p.shape, p.dtype).stop_grad()) | |||
| pg["v"].append(jt.zeros(p.shape, p.dtype).stop_grad()) | |||
| pg["d"].append(jt.zeros(p.shape, p.dtype).stop_grad()) | |||
| pg["pre_grad"].append(jt.zeros(p.shape, p.dtype).stop_grad()) | |||
| def add_param_group(self, group): | |||
| group["m"] = [] | |||
| group["v"] = [] | |||
| group["d"] = [] | |||
| group["pre_grad"] = [] | |||
| for p in group["params"]: | |||
| group["m"].append(jt.zeros(p.shape, p.dtype).stop_grad()) | |||
| group["v"].append(jt.zeros(p.shape, p.dtype).stop_grad()) | |||
| group["d"].append(jt.zeros(p.shape, p.dtype).stop_grad()) | |||
| group["pre_grad"].append(jt.zeros(p.shape, p.dtype).stop_grad()) | |||
| self.param_groups.append(group) | |||
| def step(self, loss=None, retain_graph=False): | |||
| self.pre_step(loss, retain_graph) | |||
| n = float(self.n_step) | |||
| for pg in self.param_groups: | |||
| lr = pg.get("lr", self.lr) | |||
| betas = pg.get("betas", self.betas) | |||
| eps = pg.get("eps", self.eps) | |||
| weight_decay = pg.get("weight_decay", self.weight_decay) | |||
| max_grad_norm = pg.get("max_grad_norm", self.max_grad_norm) | |||
| if max_grad_norm>0: self.clip_grad_norm(max_grad_norm) | |||
| beta1, beta2, beta3 = betas | |||
| bias_correction1 = 1 - beta1 ** n | |||
| bias_correction2 = 1 - beta2 ** n | |||
| bias_correction3_sqrt = jt.sqrt(1 - beta3 ** n) | |||
| step_size_diff = lr * beta2 * bias_correction3_sqrt / bias_correction2 | |||
| step_size = lr * bias_correction3_sqrt / bias_correction1 | |||
| eps_bias_sqrt = eps * bias_correction3_sqrt | |||
| for p, g, m, v, d, pre_g in zip(pg["params"], | |||
| pg["grads"], | |||
| pg["m"], | |||
| pg["v"], | |||
| pg["d"], | |||
| pg["pre_grad"]): | |||
| if p.is_stop_grad(): continue | |||
| if self.n_step>0: | |||
| pre_g.update(g - pre_g) # Update pre_g as grad_diff | |||
| m.update(beta1 * m + (1 - beta1) * g) | |||
| d.update(beta2 * d + (1 - beta2) * pre_g) # # Use belief as update | |||
| pre_g.update(jt.multiply(pre_g, beta2) + g) # Update pre_g as update (g + beta2 * grad_diff) | |||
| v.update(beta3 * v + (1 - beta3) * (pre_g - m) * (pre_g - m)) # Use pre_g as update | |||
| p.update(p - (step_size * m + step_size_diff * d) / (jt.sqrt(v) + eps_bias_sqrt)) | |||
| p.update(p / (1 + lr * weight_decay)) #AdanBelief best result 0.7358(300 epoch basic) | |||
| pre_g.update(g) # Update pre_g for the next iteration | |||
| self.post_step() | |||
| class LRScheduler: | |||
| def __init__(self,optimizer, last_epoch=-1): | |||
| assert isinstance(optimizer,Optimizer) | |||
| self.optimizer = optimizer | |||
| if last_epoch==-1: | |||
| for gp in optimizer.param_groups: | |||
| gp.setdefault('initial_lr',gp.get('lr',optimizer.lr)) | |||
| else: | |||
| for gp in optimizer.param_groups: | |||
| assert 'initial_lr' in gp | |||
| self.base_lrs = list(map(lambda group: group['initial_lr'], optimizer.param_groups)) | |||
| self.last_epoch = last_epoch | |||
| self.optimizer._step_count = 0 | |||
| self._step_count = 0 | |||
| self.step() | |||
| def get_lr(self): | |||
| raise NotImplementedError | |||
| def get_last_lr(self): | |||
| return self._last_lr | |||
| def step(self,epoch=None): | |||
| self._step_count += 1 | |||
| if epoch is None: | |||
| self.last_epoch += 1 | |||
| values = self.get_lr() | |||
| else: | |||
| self.last_epoch = epoch | |||
| values = self.get_lr() | |||
| for i, data in enumerate(zip(self.optimizer.param_groups, values)): | |||
| param_group, lr = data | |||
| param_group['lr'] = lr | |||
| self._last_lr = [group['lr'] for group in self.optimizer.param_groups] | |||
| class LambdaLR(LRScheduler): | |||
| def __init__(self, optimizer, lr_lambda, last_epoch=-1): | |||
| if not isinstance(lr_lambda, list) and not isinstance(lr_lambda, tuple): | |||
| self.lr_lambdas = [lr_lambda] * len(optimizer.param_groups) | |||
| else: | |||
| if len(lr_lambda) != len(optimizer.param_groups): | |||
| raise ValueError("Expected {} lr_lambdas, but got {}".format(len(optimizer.param_groups), len(lr_lambda))) | |||
| self.lr_lambdas = list(lr_lambda) | |||
| super(LambdaLR, self).__init__(optimizer, last_epoch) | |||
| def get_lr(self): | |||
| return [base_lr * lmbda(self.last_epoch) | |||
| for lmbda, base_lr in zip(self.lr_lambdas, self.base_lrs)] | |||