class_docstr: Tristage learning rate schedulr
  
  Implement the learning rate scheduler in https://arxiv.org/pdf/1904.08779.pdf
  
  Similar to inverse_squre_root scheduler, but tri_stage learning rate employs
  three stages LR scheduling:
  
      - warmup stage, starting from `lr` * `init_lr_scale`, linearly
        increased to `lr` in `warmup_steps` iterations
  
      - hold stage, after `warmup_steps`, keep the LR as `lr` for `hold_steps`
        iterations
  
      - decay stage, after hold stage, decay LR exponetially to
        `lr` * `final_lr_scale` in `decay_steps`;
        after that LR is keep as `final_lr_scale` * `lr`
  
  During warmup::
  
    init_lr = cfg.init_lr_scale * cfg.lr
    lrs = torch.linspace(init_lr, cfg.lr, cfg.warmup_steps)
    lr = lrs[update_num]
  
  During hold::
  
    lr = cfg.lr
  
  During decay::
  
    decay_factor = - math.log(cfg.final_lr_scale) / cfg.decay_steps
    lr = cfg.lr * exp(- (update_num - warmup_steps - decay_steps) * decay_factor)
  
  After that::
  
    lr = cfg.lr * cfg.final_lr_scale
code: def __init__(self, cfg: TriStageLRScheduleConfig, optimizer):
      super().__init__(cfg, optimizer)
      if len(cfg.lr) > 1:
          raise ValueError(
              "Cannot use a fixed learning rate schedule with tri-stage lr."
              " Consider --lr-scheduler=fixed instead."
          )
  
      # calculate LR at each point
      self.peak_lr = cfg.lr[0]
      self.init_lr = cfg.init_lr_scale * cfg.lr[0]
      self.final_lr = cfg.final_lr_scale * cfg.lr[0]
  
      if cfg.phase_ratio is not None:
          assert cfg.max_update > 0
          assert (
              sum(cfg.phase_ratio) == 1
          ), "phase ratios must add up to 1"
          self.warmup_steps = int(cfg.max_update * cfg.phase_ratio[0])
          self.hold_steps = int(cfg.max_update * cfg.phase_ratio[1])
          self.decay_steps = int(cfg.max_update * cfg.phase_ratio[2])
      else:
          self.warmup_steps = cfg.warmup_steps
          self.hold_steps = cfg.hold_steps
          self.decay_steps = cfg.decay_steps
  
      assert (
          self.warmup_steps + self.hold_steps + self.decay_steps > 0
      ), "please specify steps or phase_ratio"
  
      self.warmup_rate = (
          (self.peak_lr - self.init_lr) / self.warmup_steps
          if self.warmup_steps != 0
          else 0
      )
      self.decay_factor = (
          -math.log(cfg.final_lr_scale) / self.decay_steps
      )
  
      # initial learning rate
      self.lr = self.init_lr
      self.optimizer.set_lr(self.lr)
  
file_path: fairseq/fairseq/optim/lr_scheduler/tri_stage_lr_scheduler.py
imports: import math
  from dataclasses import dataclass, field
  from typing import Optional, List, Tuple
  from omegaconf import II
  from fairseq.dataclass import FairseqDataclass
  from fairseq.optim.lr_scheduler import FairseqLRScheduler, register_lr_scheduler
label_desc: Use image node_id 1 to create a new TriStageLRSchedule object from inherited base classes: FairseqLRScheduler with example: obj = TriStageLRSchedule(cfg, optimizer)