diff --git a/src/Ranger21.py b/src/Ranger21.py index 5593c40..a30d3c3 100644 --- a/src/Ranger21.py +++ b/src/Ranger21.py @@ -12,7 +12,7 @@ import torch -import torch.optim +import torch.optim as TO import math import collections @@ -20,65 +20,245 @@ import collections import copy from torch import linalg as LA -class Ranger21(torch.optim.Optimizer): - def __init__(self, - params, - lr, - eps=1e-8, - num_batches_per_epoch = None, - num_epochs = None, - num_warmup_iterations = 1000, - weight_decay=0, - decay_type = "stable", - warmup_type = 'linear', - use_GC=True): +class Ranger21(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, + num_warmup_iterations=1000, + weight_decay=0, + decay_type="stable", + warmup_type="linear", + use_GC=True, + ): - # todo - checks on incoming params - defaults = dict(lr=lr, eps=eps, weight_decay = weight_decay) - super().__init__(params, defaults) + # 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.num_warmup_iters = num_warmup_iterations - self.warmup_type=warmup_type - self.use_GC = use_GC - self.starting_lr = lr + self.num_batches = num_batches_per_epoch + self.num_epochs = num_epochs + self.num_warmup_iters = num_warmup_iterations + self.warmup_type = warmup_type + self.use_GC = use_GC + self.starting_lr = lr - #decay - self.decay = weight_decay - self.decay_type = decay_type + # decay + self.decay = weight_decay + self.decay_type = decay_type + self.param_size = 0 + + # logging + self.variance_sum_tracking = [] + + def __setstate__(self, state): + super().__setstate__(state) def warmup_dampening(self, step): # not usable yet style = self.warmup_type - step +=1 + step += 1 warmup = self.num_warmup_iters if style is None: return 1.0 - if style=='linear': - return min(1.0, (step/warmup) ) + if style == "linear": + return min(1.0, (step / warmup)) - elif style=='exponential': - return 1.0 - math.exp(-step/warmup) + elif style == "exponential": + return 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"] - @torch.no_grad - def step(self, - closure = None, - passed_loss = None): + return beta1, beta2, mean_avg, variance_avg - # let's build in a loss pass through for HyperExplorer - loss = None - if closure is not None and isinstance(closure, collections.Callable): - with torch.grad(): - loss = closure() - + # @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() + # if closure is not None: + # with torch.enable_grad(): + # loss = closure() + param_size = 0 + variance_ma_sum = 0.0 + + 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() + + # Perform optimization step + grad = p.grad + + if grad.is_sparse: + raise RuntimeError("sparse matrix not supported atm") + + state = self.state[p] + + # 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 + ) + + 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"] + + 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) + + step_size = lr / bias_correction1 + + # update weights + p.addcdiv_(grad_exp_avg, denom, value=-step_size) + + return loss + + """ param_size = 0 + variance_avg_sum = 0. + + for group in self.param_groups: + for p in group['params']: + if p.grad is None: + continue + param_size += p.numel() + + # first part of optimization + grad = p.grad + + + + state = self.state[p] + + #init if needed + if len(state)==0: + print(f"initing state") + state['step']=0 + + state['mean_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format) + state['variance_avg'] = torch.zeros_like(p,memory_format = torch.preserve_format) + + # get state values + #beta1, beta2, mean_avg, variance_avg = self.get_state_values(group, state) + beta1,beta2 = group['betas'] + mean_avg = state['mean_avg'] + variance_avg = state['variance_avg'] + + #print(f"beta1= {beta1}") + + state['step'] +=1 + + bias_correction2 = 1 - beta2**variance_avg + + # bias avgs + mean_avg.mul_(beta1).add_(grad, alpha=1-beta1) + variance_avg.mul_(beta2).addcmul_(grad,grad,value = 1-beta2) + + variance_avg_hat = variance_avg / bias_correction2 + #print(f"variance-avg-hat = {variance_avg_hat}") + + variance_avg_sum += variance_avg_hat.sum() + + print(f"param size = {param_size}") + if not self.paramsize: + self.paramsize = param_size + else: + if self.paramsize != param_size: + raise ValueError("param size changed") + print(f"variance_avg_sum = {variance_avg_sum}") + variance_avg_normalized = math.sqrt(variance_avg_sum / param_size) + print(f"variance sum normalize = {variance_avg_normalized}") +""" diff --git a/src/__pycache__/Ranger21.cpython-37.pyc b/src/__pycache__/Ranger21.cpython-37.pyc new file mode 100644 index 0000000..788bebd Binary files /dev/null and b/src/__pycache__/Ranger21.cpython-37.pyc differ