text: # Copyright (c) 2021, Parallel Systems Architecture Laboratory (PARSA), EPFL & 
  # Machine Learning and Optimization Laboratory (MLO), EPFL. All rights reserved.
  #
  # Redistribution and use in source and binary forms, with or without
  # modification, are permitted provided that the following conditions are met:
  #
  # 1. Redistributions of source code must retain the above copyright notice, this
  #    list of conditions and the following disclaimer.
  #
  # 2. Redistributions in binary form must reproduce the above copyright notice,
  #    this list of conditions and the following disclaimer in the documentation
  #    and/or other materials provided with the distribution.
  #
  # 3. Neither the name of the PARSA, EPFL & MLO, EPFL
  #    nor the names of its contributors may be used to endorse or promote
  #    products derived from this software without specific prior written
  #    permission.
  #
  # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
  # DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
  # FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
  # DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
  # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
  # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
  # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  
  # -*- coding: utf-8 -*-
  from torch.optim.optimizer import Optimizer, required
  
  from cnn.utils.auxiliary import zeros_like
  from cnn.utils.lr import adjust_learning_rate_by_lars
  from bfp.bfp_optim import get_bfp_optim
  
  class SGD(Optimizer):
      r"""Implements stochastic gradient descent (optionally with momentum).
  
      Nesterov momentum is based on the formula from
      `On the importance of initialization and momentum in deep learning`__.
  
      Args:
          params (iterable): iterable of parameters to optimize or dicts defining
              parameter groups
          lr (float): learning rate
          momentum (float, optional): momentum factor (default: 0)
          weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
          dampening (float, optional): dampening for momentum (default: 0)
          nesterov (bool, optional): enables Nesterov momentum (default: False)
  
      Example:
          >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
          >>> optimizer.zero_grad()
          >>> loss_fn(model(input), target).backward()
          >>> optimizer.step()
  
      __ http://www.cs.toronto.edu/%7Ehinton/absps/momentum.pdf
  
      .. note::
          The implementation of SGD with Momentum/Nesterov subtly differs from
          Sutskever et. al. and implementations in some other frameworks.
  
          Considering the specific case of Momentum, the update can be written as
  
          .. math::
                    v = \rho * v + g \\
                    p = p - lr * v
  
          where p, g, v and :math:`\rho` denote the parameters, gradient,
          velocity, and momentum respectively.
  
          This is in contrast to Sutskever et. al. and
          other frameworks which employ an update of the form
  
          .. math::
               v = \rho * v + lr * g \\
               p = p - v
  
          The Nesterov version is analogously modified.
      """
  
      def __init__(self, params, lr=required, momentum=0, dampening=0,
                   weight_decay=0, nesterov=False, args=None):
          defaults = dict(lr=lr, momentum=momentum, dampening=dampening,
                          weight_decay=weight_decay, nesterov=nesterov,
                          args=args)
          if nesterov and (momentum <= 0 or dampening != 0):
              raise ValueError("Nesterov momentum requires a momentum and zero dampening")
          super(SGD, self).__init__(params, defaults)
  
      def __setstate__(self, state):
          super(SGD, self).__setstate__(state)
          for group in self.param_groups:
              group.setdefault('nesterov', False)
  
      def step(self, closure=None, apply_lr=True, apply_momentum=True):
          """Performs a single optimization step.
  
          Avoid to use momentum to accumulate the gradients from other workers.
  
          Arguments:
              closure (callable, optional): A closure that reevaluates the model
                  and returns the loss.
          """
          loss = None
          if closure is not None:
              loss = closure()
  
          for group in self.param_groups:
              # retrieve para.
              args = group['args']
              weight_decay = group['weight_decay']
              momentum = group['momentum']
              dampening = group['dampening']
              nesterov = group['nesterov']
  
              for p in group['params']:
                  if p.grad is None:
                      continue
                  d_p = p.grad.data
                  if weight_decay != 0:
                      d_p.add_(p.data, alpha=weight_decay)
                  if momentum != 0 and apply_momentum:
                      param_state = self.state[p]
                      if 'momentum_buffer' not in param_state:
                          buf = param_state['momentum_buffer'] = zeros_like(p.data)
                          buf.mul_(momentum).add_(d_p)
                      else:
                          buf = param_state['momentum_buffer']
                          buf.mul_(momentum).add_(d_p, alpha=1 - dampening)
                      if nesterov:
                          d_p = d_p.add(buf, alpha=momentum)
                      else:
                          d_p = buf
  
                  # get lr from LARS.
                  lr = adjust_learning_rate_by_lars(args, group['lr'], p)
  
                  if apply_lr:
                      lr_ = lr
                  else:
                      lr_ = 1.0
  
                  p.data.add_(d_p, alpha=-lr_)
  
          return loss
  
  BFPSGD = get_bfp_optim(SGD, 'SGD')
  