Using both learning rate warm up and a learning rate scheduler

Mole_m7b5: <p>I’m trying to implement both learning rate warmup and a learning rate schedule within my training loop.</p>
<p>I’m currently using <a href="https://github.com/Tony-Y/pytorch_warmup" rel="noopener nofollow ugc">this</a> for learning rate warmup, specifically  the <code>LinearWarmup()</code>. So this simply ramps up from <code>0</code> to <code>max_lr</code> over a given number of steps.</p>
<p>I’m also wanting to use <code>CosineAnnealingWarmRestarts(optimizer, T_0, T_mult)</code> as my lr scheduler.</p>
<p>The challenge is that I’m wanting to use a rather long warm up period, without using an initially high value of <code>T_0</code>. Is there a way I can the LR scheduler to become active over after X number of steps have passed? A simplified version of my code is below. It may also be beneficial to have the LR constant for a number of steps before the scheduled decay begins.</p>
<p>If my initial value of T_0 is not high enough, the LR value during the warm up period follows the pattern of the LR schedule. I’m not sure if this is something that is desirable?</p>
<pre><code class="lang-auto">import torch
import pytorch_warmup as warmup

optimizer = torch.optim.AdamW(params, lr=lr)
num_steps = len(dataloader) * num_epochs
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=10, T_mult=2)
warmup_scheduler = warmup.LinearWarmup(optimizer, warmup_period=1000)
iters = len(dataloader)
for epoch in range(1,num_epochs+1):
    for idx,  batch in enumerate(dataloader):
        optimizer.zero_grad()
        loss = ...
        loss.backward()
        optimizer.step()
        with warmup_scheduler.dampening():
            lr_scheduler.step(epoch + idx / iters)
</code></pre>

ptrblck: <aside class="quote no-group" data-username="Mole_m7b5" data-post="1" data-topic="177767">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/mole_m7b5/48/14254_2.png" class="avatar"> Mole_m7b5:</div>
<blockquote>
<p>Is there a way I can the LR scheduler to become active over after X number of steps have passed?</p>
</blockquote>
</aside>
<p>Yes, you could use the <a href="https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.SequentialLR.html#torch.optim.lr_scheduler.SequentialLR"><code>SequentialLR</code></a> which accepts multiple schedulers and uses milestones to activate them sequentially.</p>

Tony-Y: <p>I’m the maintainer for <code>pytorch_warmup</code>.</p>
<p>Please delay the LR schedule for your need:</p>
<pre><code class="lang-auto">    with warmup_scheduler.dampening():
      if warmup_scheduler.last_step + 1 &gt;= warmup_period:
        lr_scheduler.step()
</code></pre>
<p>I put <a href="https://gist.github.com/Tony-Y/03e47588bae2e2bfd69815d046d994ec" rel="noopener nofollow ugc">Example Code</a> for CosineAnnealingWarmRestarts on Github Gist:</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/f/e/fea341721ab81e98e8336de5bceac79217143fc5.png" alt="Linear Warmup + Cosine Annealing with Warm Restarts" data-base62-sha1="AkD3arPV0fofHPf9A5Fc0yChTOR" width="423" height="309"></p>
<p><a href="https://gist.github.com/Tony-Y/7438c20c33f3a56aa6351667b33e9297" rel="noopener nofollow ugc">Example Code with T_mult=2</a>:</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/8/3/83b20140c1e27bfa428e782162c3e0ea85b99a8f.png" alt="Linear Warmup + Cosine Annealing with Warm Restarts 2" data-base62-sha1="iN1Vbxi9HEme9sn7bPgV7CkcqAf" width="423" height="309"></p>