code: def _test_noam_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)
      opts = [
          Noam(model.parameters(), d_model=512, warmup=4000, lr=1),
          Noam(model.parameters(), d_model=512, warmup=8000, lr=1),
          Noam(model.parameters(), d_model=2048, warmup=2000, lr=1),
      ]
      plt.plot(
          np.arange(1, 20000),
          [
              [opt.get_lr({"step": i}, opt.defaults) for opt in opts]
              for i in range(1, 20000)
          ],
      )
      plt.legend(["512:4000", "512:8000", "2048:2000"])
      plt.title("Learning Rate")
      plt.show()
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/optimizers/noam.py
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 for calling a global function with example usage: _test_noam_lr() without return types