diff --git a/src/Ranger21.py b/src/Ranger21.py index ac9cae7..d082a13 100644 --- a/src/Ranger21.py +++ b/src/Ranger21.py @@ -2,6 +2,8 @@ # core components based on: +# AdamP: Naver + # MADGRAD: https://arxiv.org/abs/2101.11075 # warmup: https://arxiv.org/abs/1910.04209v3 @@ -13,6 +15,7 @@ import torch import torch.optim as TO +import torch.nn.functional as F import math import collections @@ -20,11 +23,12 @@ import collections import copy from torch import linalg as LA + def centralize_gradient(x, gc_conv_only=False): """credit - https://github.com/Yonghongwei/Gradient-Centralization """ - + size = len(list(x.size())) - #print(f"size = {size}") + # print(f"size = {size}") if gc_conv_only: if size > 3: @@ -34,23 +38,25 @@ def centralize_gradient(x, gc_conv_only=False): x.add_(-x.mean(dim=tuple(range(1, size)), keepdim=True)) return x + class Ranger21(TO.Optimizer): def __init__( self, params, lr, + use_madgrad=True, + using_gc=True, + gc_conv_only=False, betas=(0.9, 0.999), # temp for checking tuned warmups momentum=0.9, eps=1e-8, num_batches_per_epoch=None, num_epochs=None, - use_warmup = True, + use_warmup=True, num_warmup_iterations=None, weight_decay=1e-4, decay_type="stable", warmup_type="linear", - use_gradient_centralization=True, - gc_conv_only=False ): # todo - checks on incoming params @@ -59,12 +65,14 @@ class Ranger21(TO.Optimizer): ) super().__init__(params, defaults) + # engine + self.use_madgrad = use_madgrad self.num_batches = num_batches_per_epoch self.num_epochs = num_epochs - + self.warmup_type = warmup_type - self.use_gc = use_gradient_centralization, - self.gc_conv_only=gc_conv_only, + self.use_gc = using_gc + self.gc_conv_only = gc_conv_only self.starting_lr = lr # decay @@ -74,31 +82,38 @@ class Ranger21(TO.Optimizer): # warmup - we'll use default recommended in Ma/Yarats unless user specifies num iterations self.use_warmup = use_warmup + if num_warmup_iterations is None: - self.num_warmup_iters = math.ceil((2 / (1-betas[1]))) # default untuned linear warmup + self.num_warmup_iters = math.ceil( + (2 / (1 - betas[1])) + ) # default untuned linear warmup else: self.num_warmup_iters = num_warmup_iterations # logging self.variance_sum_tracking = [] + # display + engine = "Adam" if not self.use_madgrad else "MadGrad" + # print out initial settings to make usage easier print(f"Ranger21 optimizer ready with following settings:\n") + print(f"Core optimizer = {engine}") print(f"Learning rate of {self.starting_lr}") + if self.use_warmup: print(f"{self.warmup_type} warmup, over {self.num_warmup_iters} iterations") + if self.decay: + print(f"Stable weight decay of {self.decay}") - print(f"Stable weight decay of {self.decay}") if self.use_gc: - print(f"Gradient Centralization = On") + print(f"Gradient Centralization = On") else: print("Gradient Centralization = Off") - def __setstate__(self, state): super().__setstate__(state) - def warmup_dampening(self, lr, step): # not usable yet style = self.warmup_type @@ -127,19 +142,17 @@ class Ranger21(TO.Optimizer): # @staticmethod @torch.no_grad() - def step(self, - closure=None): + def step(self, closure=None): loss = None if closure is not None and isinstance(closure, collections.Callable): with torch.grad(): loss = closure() - param_size = 0 variance_ma_sum = 0.0 - #phase 1 - accumulate all of the variance_ma_sum to use in stable weight decay + # phase 1 - accumulate all of the variance_ma_sum to use in stable weight decay for i, group in enumerate(self.param_groups): for j, p in enumerate(group["params"]): @@ -149,13 +162,13 @@ class Ranger21(TO.Optimizer): if not self.param_size: param_size += p.numel() - grad = p.grad if grad.is_sparse: raise RuntimeError("sparse matrix not supported atm") state = self.state[p] + momentum = group["momentum"] # State initialization if len(state) == 0: @@ -176,25 +189,37 @@ class Ranger21(TO.Optimizer): grad, gc_conv_only=self.gc_conv_only, ) - #else: + # else: # grad = uncentralized_grad state["step"] += 1 beta1, beta2 = group["betas"] grad_ma = state["grad_ma"] + variance_ma = state["variance_ma"] + step = state["step"] + lr = group["lr"] + + # if self.use_warmup: + # lr = self.warmup_dampening(lr, step) bias_correction2 = 1 - beta2 ** state["step"] # update the exp averages + # if not self.use_madgrad: grad_ma.mul_(beta1).add_(grad, alpha=1 - beta1) variance_ma.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) - variance_ma_debiased = variance_ma / bias_correction2 variance_ma_sum += variance_ma_debiased.sum() + # else: #madgrad + + # should we dupe variance_ma since stable is assuming adam style variance? + + # stable wd + # variance_ma_sum += grad_sum_sq.sum() # print(f"variance hat sum = {exp_avg_sq_hat_sum}") # Calculate the sqrt of the mean of all elements in exp_avg_sq_hat @@ -212,46 +237,118 @@ class Ranger21(TO.Optimizer): # debugging self.variance_sum_tracking.append(variance_ma_sum.item()) + # stable weight decay + # if not self.use_madgrad: variance_normalized = math.sqrt(variance_ma_sum / self.param_size) + # else: + # variance_normalized = math.pow((variance_ma / self.param_size), .3333) # print(f"variance mean sqrt = {variance_normalized}") # phase 2 - apply weight decay and step for group in self.param_groups: + + step = state["step"] + + # Perform stable weight decay + decay = group["weight_decay"] + eps = group["eps"] + lr = group["lr"] + momentum = group["momentum"] + + beta1, beta2 = group["betas"] + grad_exp_avg = state["grad_ma"] + variance_ma = state["variance_ma"] + + if self.use_warmup: + lr = self.warmup_dampening(lr, step) + + # madgrad outer + ck = 1 - momentum + lamb = lr * math.pow(step, 0.5) + + if decay: + if not self.use_madgrad: + p.data.mul_(1 - decay * lr / variance_normalized) + else: + p.data.mul_(1 - decay * lamb / variance_normalized) + + # innner loop, params for p in group["params"]: if p.grad is None: continue state = self.state[p] + grad = p.grad - step = state["step"] - # Perform stable weight decay - decay = group["weight_decay"] - eps = group["eps"] - lr = group["lr"] - if self.use_warmup: - lr = self.warmup_dampening(lr,step) - #if step < 10: - # print(f"warmup dampening at step {step} = {lr} vs {group['lr']}") + if self.use_madgrad: + if "grad_sum_sq" not in state: + state["grad_sum_sq"] = torch.zeros_like(p.data).detach() + state["s"] = torch.zeros_like(p.data).detach() + if momentum != 0: + state["x0"] = torch.clone(p.data).detach() - if decay: - p.data.mul_(1 - decay * lr / variance_normalized) + if momentum != 0.0 and grad.is_sparse: + raise RuntimeError( + "momentum != 0 is not compatible with sparse gradients" + ) - beta1, beta2 = group["betas"] - grad_exp_avg = state["grad_ma"] - variance_ma = state["variance_ma"] + # centralize gradients + if self.use_gc: + grad = centralize_gradient( + grad, + gc_conv_only=self.gc_conv_only, + ) - bias_correction1 = 1 - beta1 ** step - bias_correction2 = 1 - beta2 ** step + grad_sum_sq = state["grad_sum_sq"] + s = state["s"] + if momentum == 0: + # Compute x_0 from other known quantities + rms = grad_sum_sq.pow(1 / 3).add_(eps) + x0 = p.data.addcdiv(s, rms, value=1) + else: + x0 = state["x0"] - variance_biased_ma = variance_ma / bias_correction2 + # Accumulate second moments - denom = variance_biased_ma.sqrt().add(eps) + # print(f" grad = {grad}") + # print(f"lamb = {lamb}") + # print(f"gsumsq = {grad_sum_sq}") - step_size = lr / bias_correction1 + grad_sum_sq.addcmul_(grad, grad, value=lamb) + rms = grad_sum_sq.pow(1 / 3).add_(eps) - # update weights - p.addcdiv_(grad_exp_avg, denom, value=-step_size) + # Update s + s.data.add_(grad, alpha=lamb) + + # Step + if momentum == 0: + p.data.copy_(x0.addcdiv(s, rms, value=-1)) + else: + z = x0.addcdiv(s, rms, value=-1) + + # p is a moving average of z + p.data.mul_(1 - ck).add_(z, alpha=ck) + + else: # adam core + + beta1, beta2 = group["betas"] + grad_exp_avg = state["grad_ma"] + variance_ma = state["variance_ma"] + + # grad centralization, if used, was already used in the phase 1 pass as part of grad_exp_avg and variance_ma computations ...so no need to do it again here + bias_correction1 = 1 - beta1 ** step + bias_correction2 = 1 - beta2 ** step + + variance_biased_ma = variance_ma / bias_correction2 + + denom = variance_biased_ma.sqrt().add(eps) + + step_size = lr / bias_correction1 + + # update weights + # p.data.add_(weight_mod, alpha=-step_size) + p.addcdiv_(grad_exp_avg, denom, value=-step_size) return loss diff --git a/src/Rangerabel.py b/src/Rangerabel.py new file mode 100644 index 0000000..58b225a --- /dev/null +++ b/src/Rangerabel.py @@ -0,0 +1,335 @@ +# Ranger21 - @lessw2020 + +# core components based on: + +# AdamP: Naver + +# MADGRAD: https://arxiv.org/abs/2101.11075 + +# warmup: https://arxiv.org/abs/1910.04209v3 + +# stable weight decay: https://arxiv.org/abs/2011.11152v3 + +# Gradient Centralization: https://arxiv.org/abs/2004.01461v2 + + +import torch +import torch.optim as TO +import torch.nn.functional as F + +import math +import collections + +import copy +from torch import linalg as LA + + +def centralize_gradient(x, gc_conv_only=False): + """credit - https://github.com/Yonghongwei/Gradient-Centralization """ + + size = len(list(x.size())) + # print(f"size = {size}") + + if gc_conv_only: + if size > 3: + x.add_(-x.mean(dim=tuple(range(1, size)), keepdim=True)) + else: + if size > 1: + x.add_(-x.mean(dim=tuple(range(1, size)), keepdim=True)) + return x + + +class Ranger21abel(TO.Optimizer): + def __init__( + self, + params, + lr, + betas=(0.9, 0.999), # temp for checking tuned warmups + momentum=0.9, + eps=1e-8, + num_batches_per_epoch=None, + num_epochs=None, + use_abel=True, + abel_decay_factor = .3, + use_warmup=True, + num_warmup_iterations=None, + weight_decay=1e-4, + decay_type="stable", + warmup_type="linear", + use_gradient_centralization=True, + gc_conv_only=False, + ): + + # todo - checks on incoming params + defaults = dict( + lr=lr, momentum=momentum, betas=betas, eps=eps, weight_decay=weight_decay + ) + super().__init__(params, defaults) + + self.num_batches = num_batches_per_epoch + self.num_epochs = num_epochs + + self.warmup_type = warmup_type + self.use_gc = (use_gradient_centralization,) + self.gc_conv_only = (gc_conv_only,) + self.starting_lr = lr + self.current_lr = lr + + # abel + self.use_abel = use_abel + self.weight_list=[] + self.batch_count =0 + self.epoch = 0 + self.lr_decay_factor = abel_decay_factor + self.abel_decay_end = math.ceil(self.num_epochs * .85) + self.reached_minima = False + self.pweight_accumulator = 0 + + # decay + self.decay = weight_decay + self.decay_type = decay_type + self.param_size = 0 + + # warmup - we'll use default recommended in Ma/Yarats unless user specifies num iterations + self.use_warmup = use_warmup + if num_warmup_iterations is None: + self.num_warmup_iters = math.ceil( + (2 / (1 - betas[1])) + ) # default untuned linear warmup + else: + self.num_warmup_iters = num_warmup_iterations + + # logging + self.variance_sum_tracking = [] + + + + # print out initial settings to make usage easier + print(f"Ranger21 optimizer ready with following settings:\n") + print(f"Learning rate of {self.starting_lr}") + if self.use_warmup: + print(f"{self.warmup_type} warmup, over {self.num_warmup_iters} iterations") + + print(f"Stable weight decay of {self.decay}") + if self.use_gc: + print(f"Gradient Centralization = On") + else: + print("Gradient Centralization = Off") + print(f"Num Epochs = {self.num_epochs}") + print(f"Num batches per epoch = {self.num_batches}") + + def __setstate__(self, state): + super().__setstate__(state) + + + def warmup_dampening(self, lr, step): + # not usable yet + style = self.warmup_type + warmup = self.num_warmup_iters + + if style is None: + return 1.0 + + if style == "linear": + return lr * min(1.0, (step / warmup)) + + elif style == "exponential": + return lr * (1.0 - math.exp(-step / warmup)) + else: + raise ValueError(f"warmup type {style} not implemented.") + + def get_variance(self): + return self.variance_sum_tracking + + def get_state_values(self, group, state): + beta1, beta2 = group["betas"] + mean_avg = state["mean_avg"] + variance_avg = state["variance_avg"] + + return beta1, beta2, mean_avg, variance_avg + + def abel_update(self, step_fn, weight_norm, current_lr): + ''' update lr based on abel''' + + self.pweight_accumulator += weight_norm + + + self.batch_count +=1 + #print(f"self.batch count = {self.batch_count}") + if self.batch_count == self.num_batches: + self.epoch +=1 + self.batch_count = 0 + print(f"epoch eval for epoch {self.epoch}") + + #store weights + self.weight_list.append(self.pweight_accumulator) + + print(f"total norm for epoch {self.epoch} = {weight_norm}") + #self.pweight_accumulator = 0 + + if self.batch_count !=0: + return None + #self.epoch +=1 + new_lr = current_lr + + if len(self.weight_list) < 3: + print(len(self.weight_list)) + return step_fn + + # compute weight norm delta + if (self.weight_list[-1] - self.weight_list[-2]) * (self.weight_list[-2] - self.weight_list[-3]) < 0: + if self.reached_minima: + self.reached_minima = False + new_lr *= self.lr_decay_factor + #step_fn = self.update_train_step(self.learning_rate) + else: + self.reached_minima = True + print(f"\n*****\nABEL mininum detected, new lr = {new_lr}\n***\n") + + if self.epoch == self.abel_decay_end: + new_lr *= self.lr_decay_factor + print(f"abel final decay done, new lr = {new_lr}") + return new_lr + # @staticmethod + @torch.no_grad() + def step(self, closure=None): + + loss = None + if closure is not None and isinstance(closure, collections.Callable): + with torch.grad(): + loss = closure() + + param_size = 0 + variance_ma_sum = 0.0 + weight_norm = 0 + + + # phase 1 - accumulate all of the variance_ma_sum to use in stable weight decay + + for i, group in enumerate(self.param_groups): + for j, p in enumerate(group["params"]): + if p.grad is None: + continue + + if not self.param_size: + param_size += p.numel() + + grad = p.grad + + if grad.is_sparse: + raise RuntimeError("sparse matrix not supported atm") + + state = self.state[p] + + current_weight_norm = LA.norm(p.data) + #print(f"running norm = {current_weight_norm}") + weight_norm += current_weight_norm.item() + + # State initialization + if len(state) == 0: + # print("init state") + state["step"] = 0 + # Exponential moving average of gradient values + state["grad_ma"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + # Exponential moving average of squared gradient values + state["variance_ma"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + + # centralize gradients + if self.use_gc: + grad = centralize_gradient( + grad, + gc_conv_only=self.gc_conv_only, + ) + # else: + # grad = uncentralized_grad + + state["step"] += 1 + + beta1, beta2 = group["betas"] + grad_ma = state["grad_ma"] + variance_ma = state["variance_ma"] + + bias_correction2 = 1 - beta2 ** state["step"] + + # update the exp averages + grad_ma.mul_(beta1).add_(grad, alpha=1 - beta1) + + variance_ma.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) + + variance_ma_debiased = variance_ma / bias_correction2 + + variance_ma_sum += variance_ma_debiased.sum() + + # print(f"variance hat sum = {exp_avg_sq_hat_sum}") + # Calculate the sqrt of the mean of all elements in exp_avg_sq_hat + + # we will run this first epoch only and then memoize + if not self.param_size: + self.param_size = param_size + print(f"params size saved") + print(f"total param groups = {i+1}") + print(f"total params in groups = {j+1}") + + if not self.param_size: + raise ValueError("failed to set param size") + + # debugging + self.variance_sum_tracking.append(variance_ma_sum.item()) + + variance_normalized = math.sqrt(variance_ma_sum / self.param_size) + + # print(f"variance mean sqrt = {variance_normalized}") + + # phase 2 - apply weight decay and step + for group in self.param_groups: + for p in group["params"]: + if p.grad is None: + continue + + state = self.state[p] + + step = state["step"] + + # Perform stable weight decay + decay = group["weight_decay"] + eps = group["eps"] + #lr = group["lr"] + lr = self.current_lr + + if self.use_warmup: + lr = self.warmup_dampening(lr, step) + # if step < 10: + # print(f"warmup dampening at step {step} = {lr} vs {group['lr']}") + + if decay: + p.data.mul_(1 - decay * lr / variance_normalized) + + beta1, beta2 = group["betas"] + grad_exp_avg = state["grad_ma"] + variance_ma = state["variance_ma"] + + bias_correction1 = 1 - beta1 ** step + bias_correction2 = 1 - beta2 ** step + + variance_biased_ma = variance_ma / bias_correction2 + + denom = variance_biased_ma.sqrt().add(eps) + + weight_mod = grad_exp_avg / denom + + step_size = lr / bias_correction1 + + # update weights + #p.data.add_(weight_mod, alpha=-step_size) + p.addcdiv_(grad_exp_avg, denom, value=-step_size) + + # abel step + abel_result = self.abel_update(None, weight_norm, self.current_lr) + if abel_result is not None: + self.current_lr = abel_result + + return loss diff --git a/src/__pycache__/Ranger21.cpython-37.pyc b/src/__pycache__/Ranger21.cpython-37.pyc index f314183..d9e0954 100644 Binary files a/src/__pycache__/Ranger21.cpython-37.pyc and b/src/__pycache__/Ranger21.cpython-37.pyc differ diff --git a/src/__pycache__/Rangerabel.cpython-37.pyc b/src/__pycache__/Rangerabel.cpython-37.pyc new file mode 100644 index 0000000..205fc8b Binary files /dev/null and b/src/__pycache__/Rangerabel.cpython-37.pyc differ