How to create a scheduler which increases and decreases based on a certain LR like in Bert/Roberta

username: Santosh-Gupta
cooked: <p>I’m trying to recreate the learning rate schedules in Bert/Roberta, which start with a particular optimizer with specific args, linearly increase to a certain learning rate, and then decay with a specific rate decay.</p>
<p>Say that I am trying to reproduce the Roberta pretraining, described below:</p>
<blockquote>
<p>BERT is optimized with Adam (Kingma and Ba,<br>
2015) using the following parameters: β1 = 0.9,<br>
β2 = 0.999, ǫ = 1e-6 and L2 weight decay of 0.01. The learning rate is warmed up<br>
over the first 10,000 steps to a peak value of<br>
1e-4, and then linearly decayed. BERT trains<br>
with a dropout of 0.1 on all layers and attention weights, and a GELU activation function (Hendrycks and Gimpel, 2016). Models are<br>
pretrained for S = 1,000,000 updates, with minibatches containing B = 256 sequences of maximum length T = 512 tokens.</p>
</blockquote>
<p>I would need to  start with a learning rate of 1e-6 , warmup to 1e-4 in 1000 steps, then let the weight decay continue for the rest of the training.</p>
<p>I see that there are some learning rate scheduler here,</p>
<p><a href="https://pytorch.org/docs/stable/optim.html" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/docs/stable/optim.html</a></p>
<p>But they don’t seem to have the two phases as described in the passage above, or start/stop at certain learning rates.</p>
<p>Is there another way to reproduce the roberta learningrate schedule?</p>
<p>I’m looking at the fairseq repo, but I’m having a tricky time following the code, and I’m not sure how to copy specific parts into the trainer code I am working on</p>
<p><a href="https://pytorch.org/docs/stable/optim.html" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/docs/stable/optim.html</a></p>

username: Santosh-Gupta
cooked: <p>I have looked into how Fairseq does the warmup <a href="https://github.com/pytorch/fairseq/blob/master/examples/roberta/README.pretraining.md" rel="noopener nofollow ugc">https://github.com/pytorch/fairseq/blob/master/examples/roberta/README.pretraining.md</a></p>
<p>I am having a tricky time following the code, but it looks like the warmup is handled with PolynomialDecaySchedule (since there is a ‘<span class="mention">@register_lr_scheduler</span>(“polynomial_decay”)’, and ‘polynomial_decay’ is an arg for the pretraining command) , by passing an optimizer object, and specific args</p>
<p><strong>I can deduce most of the args from the ‘Train Roberta Base’ command here</strong><br>
<strong><a href="https://github.com/pytorch/fairseq/blob/master/examples/roberta/README.pretraining.md#2-train-roberta-base" rel="noopener nofollow ugc">https://github.com/pytorch/fairseq/blob/master/examples/roberta/README.pretraining.md#2-train-roberta-base</a></strong></p>
<p>Except for power, which I am guessing will just be left to the default of 1.</p>
<p>I am also wondering the decay is ‘and then linearly decayed’ as mentioned in the Roberta paper. I see that there is already a decay happening with AdamW, since its decay rate is set at 0.01, but I don’t think that’s a linear decrease, and the paper does mention that there is a linear decrease.</p>
<p>Also, the PolynomialDecaySchedule class mentions that its supposed to be for decaying the learning rate (“Decay the LR on a fixed schedule.”). So it sounds like the class is also handling this linear decay, based on the end_learning_rate arg.</p>
<p>It seems to be that no arg is passed in for that, so it uses the default of zero. I’m not so sure of this though, I feel like it may have been mentioned in the paper, and its a little tricky to see what exactly is being passed into PolynomialDecaySchedule , because it doesn’t seem like that class is being instantiated in the code. It seems to be instantiated a different way which I am not able to figure out.</p>
<p>Maybe the end learning rate is zero after all, since that’s how it’s done in Bert, which Roberta is based on.</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/google-research/bert/blob/master/optimization.py#L36" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/google-research/bert/blob/master/optimization.py#L36" target="_blank" rel="noopener nofollow ugc">google-research/bert/blob/master/optimization.py#L36</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="26" style="counter-reset: li-counter 25 ;">
<li>"""Creates an optimizer training op."""</li><li>global_step = tf.train.get_or_create_global_step()</li><li></li><li>learning_rate = tf.constant(value=init_lr, shape=[], dtype=tf.float32)</li><li></li><li># Implements linear decay of the learning rate.</li><li>learning_rate = tf.train.polynomial_decay(</li><li>    learning_rate,</li><li>    global_step,</li><li>    num_train_steps,</li><li class="selected">    end_learning_rate=0.0,</li><li>    power=1.0,</li><li>    cycle=False)</li><li></li><li># Implements linear warmup. I.e., if global_step &lt; num_warmup_steps, the</li><li># learning rate will be `global_step/num_warmup_steps * init_lr`.</li><li>if num_warmup_steps:</li><li>  global_steps_int = tf.cast(global_step, tf.int32)</li><li>  warmup_steps_int = tf.constant(num_warmup_steps, dtype=tf.int32)</li><li></li><li>  global_steps_float = tf.cast(global_steps_int, tf.float32)</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>Which also shows power =1.</p>

username: Santosh-Gupta
cooked: <p>I am attempting to make a custom scheduler which replicates the Roberta warmup, so far I came up with this, based on Huggingface’s linear warmup scheduler</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/huggingface/transformers/blob/master/src/transformers/optimization.py#L71" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/huggingface/transformers/blob/master/src/transformers/optimization.py#L71" target="_blank" rel="noopener nofollow ugc">huggingface/transformers/blob/master/src/transformers/optimization.py#L71</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="61" style="counter-reset: li-counter 60 ;">
<li>    """</li><li></li><li>    def lr_lambda(current_step: int):</li><li>        if current_step &lt; num_warmup_steps:</li><li>            return float(current_step) / float(max(1.0, num_warmup_steps))</li><li>        return 1.0</li><li></li><li>    return LambdaLR(optimizer, lr_lambda, last_epoch=last_epoch)</li><li></li><li></li><li class="selected">def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps, last_epoch=-1):</li><li>    """</li><li>    Create a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to 0, after</li><li>    a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer.</li><li></li><li>    Args:</li><li>        optimizer (:class:`~torch.optim.Optimizer`):</li><li>            The optimizer for which to schedule the learning rate.</li><li>        num_warmup_steps (:obj:`int`):</li><li>            The number of steps for the warmup phase.</li><li>        num_training_steps (:obj:`int`):</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<pre><code class="lang-auto">def get_linear_schedule_with_warmup_with_peak(optimizer, num_warmup_steps, num_training_steps, init_lr, peak_lr, last_epoch=-1):

    def lr_lambda(current_step: int):
        if current_step &lt; num_warmup_steps:
            return (float(current_step) / float(max(1, num_warmup_steps)))*(peak_lr/init_lr)
        return max(
            0.0, float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps))
        )

    return LambdaLR(optimizer, lr_lambda, last_epoch)
</code></pre>
<p>Which is not quite exact, since the optimizer is AdamW, which has its own weight decay. For it to be exact, I need init_lr to be replaced by the current optimizer learning rate, but from the documentation on LambdaLR, lr_lambda seems to only take in an int.</p>
<aside class="onebox allowlistedgeneric">
  <header class="source">
      <a href="https://pytorch.org/docs/stable/optim.html" target="_blank" rel="noopener nofollow ugc">pytorch.org</a>
  </header>
  <article class="onebox-body">
    <img src="" class="thumbnail" width="" height="">

<h3><a href="https://pytorch.org/docs/stable/optim.html" target="_blank" rel="noopener nofollow ugc">torch.optim — PyTorch 1.7.0 documentation</a></h3>



  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>So I need a way to adjust the optimizer’s learning rate based on its current learning rate.</p>