text: #!/usr/bin/env python3
  # -*- coding: utf-8 -*-
  
  # ==============================================================================
  #
  # Copyright (c) 2020-* Less Wright. All Rights Reserved.
  #                      [orig. code: https://github.com/lessw2020/Ranger-Deep-Learning-Optimizer ]
  #
  # Copyright (c) 2020-* Emanuele Ballarin <emanuele@ballarin.cc>
  #                      All Rights Reserved.
  #                      [maintainance, adaptation, extension]
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     https://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  #
  # ==============================================================================
  #
  # SPDX-License-Identifier: Apache-2.0
  #
  # Ranger version: 2020.9.4 (ranger2020.py)
  # Cfr.: https://github.com/lessw2020/Ranger-Deep-Learning-Optimizer
  
  import math
  import torch
  from torch.optim.optimizer import Optimizer
  
  
  def centralized_gradient(x, use_gc=True, gc_conv_only=False):
      """credit - https://github.com/Yonghongwei/Gradient-Centralization"""
      if use_gc:
          if gc_conv_only:
              if len(list(x.size())) > 3:
                  x.add_(-x.mean(dim=tuple(range(1, len(list(x.size())))), keepdim=True))
          else:
              if len(list(x.size())) > 1:
                  x.add_(-x.mean(dim=tuple(range(1, len(list(x.size())))), keepdim=True))
      return x
  
  
  class Ranger(Optimizer):
      def __init__(
          self,
          params,
          lr=1e-3,  # lr
          alpha=0.5,
          k=6,
          N_sma_threshhold=5,  # Ranger options
          betas=(0.95, 0.999),
          eps=1e-5,
          weight_decay=0,  # Adam options
          # Gradient centralization on or off, applied to conv layers only or conv + fc layers
          use_gc=True,
          gc_conv_only=False,
          gc_loc=True,
      ):
  
          # parameter checks
          if not 0.0 <= alpha <= 1.0:
              raise ValueError(f"Invalid slow update rate: {alpha}")
          if k < 1:
              raise ValueError(f"Invalid lookahead steps: {k}")
          if lr <= 0:
              raise ValueError(f"Invalid Learning Rate: {lr}")
          if eps <= 0:
              raise ValueError(f"Invalid eps: {eps}")
  
          # parameter comments:
          # beta1 (momentum) of .95 seems to work better than .90...
          # N_sma_threshold of 5 seems better in testing than 4.
          # In both cases, worth testing on your dataset (.90 vs .95, 4 vs 5) to make sure which works best for you.
  
          # prep defaults and init torch.optim base
          defaults = dict(
              lr=lr,
              alpha=alpha,
              k=k,
              step_counter=0,
              betas=betas,
              N_sma_threshhold=N_sma_threshhold,
              eps=eps,
              weight_decay=weight_decay,
          )
          super().__init__(params, defaults)
  
          # adjustable threshold
          self.N_sma_threshhold = N_sma_threshhold
  
          # look ahead params
  
          self.alpha = alpha
          self.k = k
  
          # radam buffer for state
          self.radam_buffer = [[None, None, None] for _ in range(10)]
  
          # gc on or off
          self.gc_loc = gc_loc
          self.use_gc = use_gc
          self.gc_conv_only = gc_conv_only
          # level of gradient centralization
          # self.gc_gradient_threshold = 3 if gc_conv_only else 1
  
          print(
              f"Ranger optimizer loaded. \nGradient Centralization usage = {self.use_gc}"
          )
          if self.use_gc and not self.gc_conv_only:
              print("GC applied to both conv and fc layers")
          elif self.use_gc and self.gc_conv_only:
              print("GC applied to conv layers only")
  
      def __setstate__(self, state):
          print("set state called")
          super(Ranger, self).__setstate__(state)
  
      def step(self, closure=None):
          loss = None
          # note - below is commented out b/c I have other work that passes back the loss as a float, and thus not a callable closure.
          # Uncomment if you need to use the actual closure...
  
          # if closure is not None:
          # loss = closure()
  
          # Evaluate averages and grad, update param tensors
          for group in self.param_groups:
  
              for p in group["params"]:
                  if p.grad is None:
                      continue
                  grad = p.grad.data.float()
  
                  if grad.is_sparse:
                      raise RuntimeError(
                          "Ranger optimizer does not support sparse gradients"
                      )
  
                  p_data_fp32 = p.data.float()
  
                  state = self.state[p]  # get state dict for this param
  
                  if (
                      not state
                  ):  # if first time to run...init dictionary with our desired entries
                      # if self.first_run_check==0:
                      # self.first_run_check=1
                      # print("Initializing slow buffer...should not see this at load from saved model!")
                      state["step"] = 0
                      state["exp_avg"] = torch.zeros_like(p_data_fp32)
                      state["exp_avg_sq"] = torch.zeros_like(p_data_fp32)
  
                      # look ahead weight storage now in state dict
                      state["slow_buffer"] = torch.empty_like(p.data)
                      state["slow_buffer"].copy_(p.data)
  
                  else:
                      state["exp_avg"] = state["exp_avg"].type_as(p_data_fp32)
                      state["exp_avg_sq"] = state["exp_avg_sq"].type_as(p_data_fp32)
  
                  # begin computations
                  exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
                  beta1, beta2 = group["betas"]
  
                  # GC operation for Conv layers and FC layers
                  # if grad.dim() > self.gc_gradient_threshold:
                  #    grad.add_(-grad.mean(dim=tuple(range(1, grad.dim())), keepdim=True))
                  if self.gc_loc:
                      grad = centralized_gradient(
                          grad, use_gc=self.use_gc, gc_conv_only=self.gc_conv_only
                      )
  
                  state["step"] += 1
  
                  # compute variance mov avg
                  exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
  
                  # compute mean moving avg
                  exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
  
                  buffered = self.radam_buffer[int(state["step"] % 10)]
  
                  if state["step"] == buffered[0]:
                      N_sma, step_size = buffered[1], buffered[2]
                  else:
                      buffered[0] = state["step"]
                      beta2_t = beta2 ** state["step"]
                      N_sma_max = 2 / (1 - beta2) - 1
                      N_sma = N_sma_max - 2 * state["step"] * beta2_t / (1 - beta2_t)
                      buffered[1] = N_sma
                      if N_sma > self.N_sma_threshhold:
                          step_size = math.sqrt(
                              (1 - beta2_t)
                              * (N_sma - 4)
                              / (N_sma_max - 4)
                              * (N_sma - 2)
                              / N_sma
                              * N_sma_max
                              / (N_sma_max - 2)
                          ) / (1 - beta1 ** state["step"])
                      else:
                          step_size = 1.0 / (1 - beta1 ** state["step"])
                      buffered[2] = step_size
  
                  # if group['weight_decay'] != 0:
                  #    p_data_fp32.add_(-group['weight_decay']
                  #                     * group['lr'], p_data_fp32)
  
                  # apply lr
                  if N_sma > self.N_sma_threshhold:
                      denom = exp_avg_sq.sqrt().add_(group["eps"])
                      G_grad = exp_avg / denom
                  else:
                      G_grad = exp_avg
  
                  if group["weight_decay"] != 0:
                      G_grad.add_(p_data_fp32, alpha=group["weight_decay"])
                  # GC operation
                  if not self.gc_loc:
                      G_grad = centralized_gradient(
                          G_grad, use_gc=self.use_gc, gc_conv_only=self.gc_conv_only
                      )
  
                  p_data_fp32.add_(G_grad, alpha=-step_size * group["lr"])
                  p.data.copy_(p_data_fp32)
  
                  # integrated look ahead...
                  # we do it at the param level instead of group level
                  if state["step"] % group["k"] == 0:
                      # get access to slow param tensor
                      slow_p = state["slow_buffer"]
                      # (fast weights - slow weights) * alpha
                      slow_p.add_(p.data - slow_p, alpha=self.alpha)
                      # copy interpolated weights to RAdam param tensor
                      p.data.copy_(slow_p)
  
          return loss
  