class_docstr: ## Adam Optimizer with Warmup
  
  This class extends from AMSGrad optimizer defined in [`amsgrad.py`](amsgrad.html).
class_docstr_tok: ["#","#","Adam","Optimizer","with","Warmup","This","class","extends","from","AMSGrad","optimizer","defined","in","[","`","amsgrad.py","`","]","(","amsgrad.html",")","."]
docstr: ### Initialize the optimizer
  
  * `params` is the list of parameters
  * `lr` is the learning rate $lpha$
  * `betas` is a tuple of ($eta_1$, $eta_2$)
  * `eps` is $\hat{\epsilon}$ or $\epsilon$ based on `optimized_update`
  * `weight_decay` is an instance of class `WeightDecay` defined in [`__init__.py`](index.html)
  * 'optimized_update' is a flag whether to optimize the bias correction of the second moment
    by doing it after adding $\epsilon$
  * `amsgrad` is a flag indicating whether to use AMSGrad or fallback to plain Adam
  * `warmup` number of warmup steps
  * `defaults` is a dictionary of default for group values.
   This is useful when you want to extend the class `AdamWarmup`.
docstr_tok: ["#","#","#","Initialize","the","optimizer","*","`","params","`","is","the","list","of","parameters","*","`","lr","`","is","the","learning","rate","$","\u0007lpha","$","*","`","betas","`","is","a","tuple","of","(","$","\beta_1","$",",","$","\beta_2","$",")","*","`","eps","`","is","$","\\hat","{","\\epsilon","}","$","or","$","\\epsilon","$","based","on","`","optimized_update","`","*","`","weight_decay","`","is","an","instance","of","class","`","WeightDecay","`","defined","in","[","`","__init__.py","`","]","(","index.html",")","*","'optimized_update","'","is","a","flag","whether","to","optimize","the","bias","correction","of","the","second","moment","by","doing","it","after","adding","$","\\epsilon","$","*","`","amsgrad","`","is","a","flag","indicating","whether","to","use","AMSGrad","or","fallback","to","plain","Adam","*","`","warmup","`","number","of","warmup","steps","*","`","defaults","`","is","a","dictionary","of","default","for","group","values",".","This","is","useful","when","you","want","to","extend","the","class","`","AdamWarmup","`","."]
code: def __init__(
      self,
      params,
      lr=1e-3,
      betas=(0.9, 0.999),
      eps=1e-16,
      weight_decay: WeightDecay = WeightDecay(),
      optimized_update: bool = True,
      amsgrad=False,
      warmup=0,
      defaults=None,
  ):
      """
      ### Initialize the optimizer
  
      * `params` is the list of parameters
      * `lr` is the learning rate $\alpha$
      * `betas` is a tuple of ($\beta_1$, $\beta_2$)
      * `eps` is $\hat{\epsilon}$ or $\epsilon$ based on `optimized_update`
      * `weight_decay` is an instance of class `WeightDecay` defined in [`__init__.py`](index.html)
      * 'optimized_update' is a flag whether to optimize the bias correction of the second moment
        by doing it after adding $\epsilon$
      * `amsgrad` is a flag indicating whether to use AMSGrad or fallback to plain Adam
      * `warmup` number of warmup steps
      * `defaults` is a dictionary of default for group values.
       This is useful when you want to extend the class `AdamWarmup`.
      """
  
      defaults = {} if defaults is None else defaults
      defaults.update(dict(warmup=warmup))
      super().__init__(
          params,
          lr,
          betas,
          eps,
          weight_decay,
          optimized_update,
          amsgrad,
          defaults,
      )
  
imports: from typing import Dict
  from labml_nn.optimizers import WeightDecay
  from labml_nn.optimizers.amsgrad import AMSGrad
label_desc: Use image node_id 1 to create a new AdamWarmup object from inherited base classes: AMSGrad with example: obj = AdamWarmup(params, lr, betas, eps, weight_decay, optimized_update, amsgrad, warmup, defaults)