code: def _test_lr():
      """
      ### Plot learning rate for different warmups and model sizes
  
      ![Plot of learning rate](noam_lr.png)
      """
      import matplotlib.pyplot as plt
      import numpy as np
      from torch import nn
  
      model = nn.Linear(10, 10)
      opt = AdamWarmupCosineDecay(
          model.parameters(), warmup=5000, lr=1e-4, total_steps=4e6
      )
      steps = 20_000
      plt.plot(
          np.arange(1, steps),
          [
              opt.get_lr({"step": i}, opt.defaults)
              for i in range(1, steps)
          ],
      )
      plt.legend(["5000:4e6", "5000:2e6", "5000:1e6"])
      plt.title("Learning Rate")
      plt.show()
  
      steps = int(6e6)
      step_size = 1000
      plt.plot(
          np.arange(1, steps, step_size),
          [
              opt.get_lr({"step": i}, opt.defaults)
              for i in range(1, steps, step_size)
          ],
      )
      plt.legend(["5000:4e6", "5000:2e6", "5000:1e6"])
      plt.title("Learning Rate")
      plt.show()
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/optimizers/adam_warmup_cosine_decay.py
imports: import math
  from typing import Dict
  from labml_nn.optimizers import WeightDecay
  from labml_nn.optimizers.amsgrad import AMSGrad
label_desc: Use image node_id 1 for calling a global function with example usage: _test_lr() without return types