Skip to content

Instantly share code, notes, and snippets.

@Deepayan137

Deepayan137/stlr.py

Created February 8, 2020 06:23
Show Gist options
  • Save Deepayan137/16675f7039b83be1276dc6dbaa674a82 to your computer and use it in GitHub Desktop.
Save Deepayan137/16675f7039b83be1276dc6dbaa674a82 to your computer and use it in GitHub Desktop.
Download ZIP

Revisions

  1. Deepayan137 created this gist Feb 8, 2020.
    16 changes: 16 additions & 0 deletions stlr.py
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,16 @@
    'STLR scheduler from https://arxiv.org/abs/1801.06146'

    class STLR(_LRScheduler):
    def __init__(self, optimizer, T_max, last_epoch=-1, ratio=32):
    self.T_max = T_max
    self.cut = np.floor(T_max*0.1)
    self.ratio = ratio
    super(STLR, self).__init__(optimizer, last_epoch)

    def get_lr(self):
    if self.last_epoch < self.cut:
    p = self.last_epoch/self.cut
    else:
    fraction = (self.last_epoch - self.cut)/(self.cut*(1/self.ratio - 1))
    p = 1 - fraction
    return [base_lr*(1 + p*(self.ratio - 1))/self.ratio for base_lr in self.base_lrs]