class_docstr: ## Adam Optimizer
  
  We extend the class `GenericAdaptiveOptimizer` defined in [`__init__.py`](index.html)
  to implement the Adam optimizer.
docstr: ### Do the *Adam* parameter update
  
  * `state` is the optimizer state of the parameter (tensor)
  * `group` stores optimizer attributes of the parameter group
  * `param` is the parameter tensor $     heta_{t-1}$
  * `m` and `v` are the uncorrected first and second moments $m_t$ and $v_t$.
  
  This computes the following
  
  egin{align}
          heta_t &\leftarrow      heta_{t-1} - lpha \cdot rac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}
  \end{align}
  
  Since $lpha$, $eta_1$, $eta_2$ and $\epsilon$ are scalars and others are tensors
  we modify this calculation to optimize the computation.
  
  egin{align}
          heta_t &\leftarrow      heta_{t-1} - lpha \cdot rac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} \
          heta_t &\leftarrow      heta_{t-1} - lpha \cdot
          rac{m_t / (1-eta_1^t)}{\sqrt{v_t/(1-eta_2^t)} + \epsilon} \
          heta_t &\leftarrow      heta_{t-1} - lpha rac{\sqrt{1-eta_2^t}}{1-eta_1^t} \cdot
          rac{m_t}{\sqrt{v_t} + \hat{\epsilon}} \
  \end{align}
  
  where
  $$\hat{\epsilon} = (1-eta_2^t) \epsilon$$
  is what we should specify as the hyper-parameter.
docstr_tok: ["#","#","#","Do","the","*","Adam","*","parameter","update","*","`","state","`","is","the","optimizer","state","of","the","parameter","(","tensor",")","*","`","group","`","stores","optimizer","attributes","of","the","parameter","group","*","`","param","`","is","the","parameter","tensor","$","heta_","{","t-1","}","$","*","`","m","`","and","`","v","`","are","the","uncorrected","first","and","second","moments","$","m_t","$","and","$","v_t","$",".","This","computes","the","following","\begin","{","align","}","heta_t","&","\\leftarrow","heta_","{","t-1","}","-","\u0007lpha","\\cdot","rac","{","\\hat","{","m","}","_t","}","{","\\sqrt","{","\\hat","{","v","}","_t","}","+","\\epsilon","}","\\end","{","align","}","Since","$","\u0007lpha","$",",","$","\beta_1","$",",","$","\beta_2","$","and","$","\\epsilon","$","are","scalars","and","others","are","tensors","we","modify","this","calculation","to","optimize","the","computation",".","\begin","{","align","}","heta_t","&","\\leftarrow","heta_","{","t-1","}","-","\u0007lpha","\\cdot","rac","{","\\hat","{","m","}","_t","}","{","\\sqrt","{","\\hat","{","v","}","_t","}","+","\\epsilon","}","\\","heta_t","&","\\leftarrow","heta_","{","t-1","}","-","\u0007lpha","\\cdot","rac","{","m_t","\/","(","1-\beta_1^t",")","}","{","\\sqrt","{","v_t\/","(","1-\beta_2^t",")","}","+","\\epsilon","}","\\","heta_t","&","\\leftarrow","heta_","{","t-1","}","-","\u0007lpha","rac","{","\\sqrt","{","1-\beta_2^t","}","}","{","1-\beta_1^t","}","\\cdot","rac","{","m_t","}","{","\\sqrt","{","v_t","}","+","\\hat","{","\\epsilon","}","}","\\","\\end","{","align","}","where","$","$","\\hat","{","\\epsilon","}","=","(","1-\beta_2^t",")","\\epsilon","$","$","is","what","we","should","specify","as","the","hyper-parameter","."]
code: def adam_update(
      self,
      state: Dict[str, any],
      group: Dict[str, any],
      param: torch.nn.Parameter,
      m: torch.Tensor,
      v: torch.Tensor,
  ):
      """
          ### Do the *Adam* parameter update
  
          * `state` is the optimizer state of the parameter (tensor)
          * `group` stores optimizer attributes of the parameter group
          * `param` is the parameter tensor $\theta_{t-1}$
          * `m` and `v` are the uncorrected first and second moments $m_t$ and $v_t$.
  
          This computes the following
  
          \begin{align}
          \theta_t &\leftarrow \theta_{t-1} - \alpha \cdot \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}
          \end{align}
  
          Since $\alpha$, $\beta_1$, $\beta_2$ and $\epsilon$ are scalars and others are tensors
          we modify this calculation to optimize the computation.
  
          \begin{align}
          \theta_t &\leftarrow \theta_{t-1} - \alpha \cdot \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} \\
          \theta_t &\leftarrow \theta_{t-1} - \alpha \cdot
                  \frac{m_t / (1-\beta_1^t)}{\sqrt{v_t/(1-\beta_2^t)} + \epsilon} \\
          \theta_t &\leftarrow \theta_{t-1} - \alpha \frac{\sqrt{1-\beta_2^t}}{1-\beta_1^t} \cdot
                  \frac{m_t}{\sqrt{v_t} + \hat{\epsilon}} \\
          \end{align}
  
          where
          $$\hat{\epsilon} = (1-\beta_2^t) \epsilon$$
          is what we should specify as the hyper-parameter.
          """
  
      # Get $\beta_1$ and $\beta_2$
      beta1, beta2 = group["betas"]
      # Bias correction term for $\hat{m}_t$, $1 - \beta_1^t$
      bias_correction1 = 1 - beta1 ** state["step"]
      # Bias correction term for $\hat{v}_t$, $1 - \beta_2^t$
      bias_correction2 = 1 - beta2 ** state["step"]
  
      # Get learning rate
      lr = self.get_lr(state, group)
  
      # Whether to optimize the computation
      if self.optimized_update:
          # $\sqrt{v_t} + \hat{\epsilon}$
          denominator = v.sqrt().add_(group["eps"])
          # $\alpha \frac{\sqrt{1-\beta_2^t}}{1-\beta_1^t}$
          step_size = (
              lr * math.sqrt(bias_correction2) / bias_correction1
          )
          # $\theta_t \leftarrow \theta_{t-1} - \alpha \frac{\sqrt{1-\beta_2^t}}{1-\beta_1^t} \cdot
          #  \frac{m_t}{\sqrt{v_t} + \hat{\epsilon}}$
          param.data.addcdiv_(m, denominator, value=-step_size)
      # Computation without optimization
      else:
          # $\frac{\sqrt{v_t}}{\sqrt{1-\beta_2^t}} + \epsilon$
          denominator = (v.sqrt() / math.sqrt(bias_correction2)).add_(
              group["eps"]
          )
          # $\frac{\alpha}{1-\beta_1^t}$
          step_size = lr / bias_correction1
          # $\theta_t \leftarrow \theta_{t-1} - \alpha \cdot
          # \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}$
          param.data.addcdiv_(m, denominator, value=-step_size)
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/optimizers/adam.py
imports: import math
  from typing import Dict, Any, Tuple, Optional
  import torch
  from labml import tracker
  from torch import nn
  from labml_nn.optimizers import GenericAdaptiveOptimizer, WeightDecay
label_desc: Use image node_id 5 for calling the Adam obj's underlying member method code with example usage: obj.adam_update(state, group, param, m, v) without return types