class_docstr: Implements NovoGrad algorithm.
  
  Currently GPU-only.  Requires Apex to be installed via
  ``pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./``.
  
  This version of fused NovoGrad implements 2 fusions.
  
    * Fusion of the NovoGrad update's elementwise operations
    * A multi-tensor apply launch that batches the elementwise updates applied to all the model's parameters into one or a few kernel launches.
  
  :class:`apex.optimizers.FusedNovoGrad`'s usage is identical to any Pytorch optimizer::
  
      opt = apex.optimizers.FusedNovoGrad(model.parameters(), lr = ....)
      ...
      opt.step()
  
  :class:`apex.optimizers.FusedNovoGrad` may be used with or without Amp.  If you wish to use :class:`FusedNovoGrad` with Amp,
  you may choose any ``opt_level``::
  
      opt = apex.optimizers.FusedNovoGrad(model.parameters(), lr = ....)
      model, opt = amp.initialize(model, opt, opt_level="O0" or "O1 or "O2")
      ...
      opt.step()
  
  In general, ``opt_level="O1"`` is recommended.
  
  It has been proposed in `Jasper: An End-to-End Convolutional Neural Acoustic Model`_.
  More info: https://nvidia.github.io/OpenSeq2Seq/html/optimizers.html#novograd
  
  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], optional): coefficients used for computing
          running averages of gradient and its norm. (default: (0.9, 0.999))
      eps (float, optional): term added to the denominator to improve
          numerical stability. (default: 1e-8)
      weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
      amsgrad (boolean, optional): whether to use the AMSGrad variant of this
          algorithm from the paper `On the Convergence of Adam and Beyond`_
          NOT SUPPORTED now! (default: False)
      reg_inside_moment (bool, optional): whether do regularization (norm and L2)
          in momentum calculation. True for include, False for not include and
          only do it on update term. (default: False)
      grad_averaging (bool, optional): whether apply (1-beta1) to grad when
          calculating running averages of gradient. (default: True)
      norm_type (int, optional): which norm to calculate for each layer.
          2 for L2 norm, and 0 for infinite norm. These 2 are only supported
          type now. (default: 2)
      init_zero (bool, optional): whether init norm with 0 (start averaging on
          1st step) or first step norm (start averaging on 2nd step). True for
          init with 0. (default: False)
      set_grad_none (bool, optional): whether set grad to None when zero_grad()
          method is called. (default: True)
  
  .. _Jasper - An End-to-End Convolutional Neural Acoustic Model:
      https://arxiv.org/abs/1904.03288
  .. _On the Convergence of Adam and Beyond:
      https://openreview.net/forum?id=ryQu7f-RZ
docstr: Performs a single optimization step.
  
  Arguments:
      closure (callable, optional): A closure that reevaluates the model
          and returns the loss.
docstr_tok: ["Performs","a","single","optimization","step",".","Arguments",":","closure","(","callable",",","optional",")",":","A","closure","that","reevaluates","the","model","and","returns","the","loss","."]
code: def step(self, closure=None):
      """Performs a single optimization step.
  
      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:
          bias_correction = 1 if group["bias_correction"] else 0
          beta1, beta2 = group["betas"]
          grad_averaging = 1 if group["grad_averaging"] else 0
  
          # assume same step across group now to simplify things
          # per parameter step can be easily support by making it tensor, or pass list into kernel
          if "step" in group:
              group["step"] += 1
          else:
              group["step"] = 1
  
          # create lists for multi-tensor apply
          g_16, p_16, m_16 = [], [], []
          g_32, p_32, m_32 = [], [], []
  
          for p in group["params"]:
              if p.grad is None:
                  continue
              if p.grad.data.is_sparse:
                  raise RuntimeError(
                      "FusedNovoGrad does not support sparse gradients, please consider SparseAdam instead"
                  )
  
              state = self.state[p]
              # State initialization
              if len(state) == 0:
                  # Exponential moving average of gradient values
                  state["exp_avg"] = torch.zeros_like(p.data)
  
              if p.dtype == torch.float16:
                  g_16.append(p.grad.data)
                  p_16.append(p.data)
                  m_16.append(state["exp_avg"])
              elif p.dtype == torch.float32:
                  g_32.append(p.grad.data)
                  p_32.append(p.data)
                  m_32.append(state["exp_avg"])
              else:
                  raise RuntimeError(
                      "FusedNovoGrad only support fp16 and fp32."
                  )
  
          # we store per weight norm as one tensor for one group/precision combination
          # different from optim.Adam, we store norm here(not ^2) so we can unify calculation for norm types
          if "exp_avg_sq" not in group:
              group["exp_avg_sq"] = [None, None]
              if group["init_zero"]:
                  # Creating the following parameters on the same device as the params tensors.
                  group["exp_avg_sq"][0] = (
                      torch.cuda.FloatTensor(
                          len(g_16),
                          device=self.param_groups[0]["params"][
                              0
                          ].device,
                      )
                      .contiguous()
                      .fill_(0)
                  )
                  group["exp_avg_sq"][1] = (
                      torch.cuda.FloatTensor(
                          len(g_32),
                          device=self.param_groups[0]["params"][
                              0
                          ].device,
                      )
                      .contiguous()
                      .fill_(0)
                  )
              else:  # init with first step norm, so first blend have no effect
                  if group["norm_type"] == 0:
                      v_16 = [
                          torch.max(
                              torch.abs(g.to(torch.float32))
                          ).item()
                          for g in g_16
                      ]
                      v_32 = [
                          torch.max(torch.abs(g)).item() for g in g_32
                      ]
                  elif group["norm_type"] == 2:
                      v_16 = [
                          torch.sum(torch.pow(g.to(torch.float32), 2))
                          .sqrt()
                          .item()
                          for g in g_16
                      ]
                      v_32 = [
                          torch.sum(torch.pow(g, 2)).sqrt().item()
                          for g in g_32
                      ]
                  else:
                      raise RuntimeError(
                          "FusedNovoGrad only support l2/inf norm now."
                      )
                  # Creating the following parameters on the same device as the params tensors.
                  group["exp_avg_sq"][0] = torch.cuda.FloatTensor(
                      v_16,
                      device=self.param_groups[0]["params"][0].device,
                  )
                  group["exp_avg_sq"][1] = torch.cuda.FloatTensor(
                      v_32,
                      device=self.param_groups[0]["params"][0].device,
                  )
          else:
              assert len(g_16) == group["exp_avg_sq"][0].numel()
              assert len(g_32) == group["exp_avg_sq"][1].numel()
  
          if len(g_16) > 0:
              multi_tensor_applier(
                  self.multi_tensor_novograd,
                  self._dummy_overflow_buf,
                  [g_16, p_16, m_16],
                  group["exp_avg_sq"][0],
                  group["lr"],
                  beta1,
                  beta2,
                  group["eps"],
                  group["step"],
                  bias_correction,
                  group["weight_decay"],
                  grad_averaging,
                  self.moment_mode,
                  group["norm_type"],
              )
          if len(g_32) > 0:
              multi_tensor_applier(
                  self.multi_tensor_novograd,
                  self._dummy_overflow_buf,
                  [g_32, p_32, m_32],
                  group["exp_avg_sq"][1],
                  group["lr"],
                  beta1,
                  beta2,
                  group["eps"],
                  group["step"],
                  bias_correction,
                  group["weight_decay"],
                  grad_averaging,
                  self.moment_mode,
                  group["norm_type"],
              )
  
      return loss
  
file_path: apex/apex/optimizers/fused_novograd.py
imports: import torch
  from apex.multi_tensor_apply import multi_tensor_applier
label_desc: Use image node_id 4 for calling the FusedNovoGrad obj's underlying member method code with example usage: obj.step(closure) and returns: loss