LR warm-up phase using SequentialLR with ExponentialLR and CosineLR

Laurent1: <p>Hi,</p>
<p>I would like to create a learning rate warm-up phase using <code>SequentialLR</code> to transition from <code>ExponentialLR</code> to <code>CosineAnnealingLR</code>.</p>
<p>Why? When going from <code>LinearLR</code> to <code>CosineAnnealingLR</code>, the learning rates are essentially close to the max during the warm-up phase. Here is my code for this case:</p>
<pre><code class="lang-auto">from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR, ExponentialLR
from torch import optim
from torch import nn
import matplotlib.pyplot as plt
import numpy as np
    
max_lr = 0.01
initial_lr = max_lr / 1e4
model = nn.Linear(10,10)
opt = optim.SGD(model.parameters(), lr=max_lr, momentum=0.9)
num_epochs = 20
warmup_scheduler = LinearLR(opt,
                         start_factor=1/1e4,
                         end_factor=1,
                         total_iters=8)
cos_scheduler = CosineAnnealingLR(opt,
                                  T_max = num_epochs-8,
                                  eta_min=max_lr / 10000.)
scheduler = SequentialLR(opt,
                         schedulers=[warmup_scheduler, cos_scheduler],
                         milestones=[8])
lrs = []
for epoch in range(num_epochs):
    lr = scheduler.get_last_lr()[0]
    lrs.append(lr)
    print(f"epoch: {epoch}: lr = {lr}")
    scheduler.step()
    
fig = plt.figure(figsize=(8,4))
ax = fig.add_subplot(1,1,1)
ax.plot(np.arange(len(lrs)),lrs)
# ax.set_yscale('log')
fig.tight_layout()
plt.savefig('foo.png', bbox_inches='tight')
plt.show()
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/b/3b5042f8bfd35d683d78f35c84a6173a31f513c1.png" data-download-href="https://discuss.pytorch.org/uploads/default/3b5042f8bfd35d683d78f35c84a6173a31f513c1" title="linear"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/3/b/3b5042f8bfd35d683d78f35c84a6173a31f513c1_2_690x345.png" alt="linear" data-base62-sha1="8sI85tZaH7Wo55Z2KyLWjLkaoLL" width="690" height="345" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/3/b/3b5042f8bfd35d683d78f35c84a6173a31f513c1_2_690x345.png, https://discuss.pytorch.org/uploads/default/original/3X/3/b/3b5042f8bfd35d683d78f35c84a6173a31f513c1.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/3/b/3b5042f8bfd35d683d78f35c84a6173a31f513c1.png 2x" data-dominant-color="FBFCFC"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">linear</span><span class="informations">800×400 17.5 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>When replacing <code>LinearLR</code> by <code>ExponentialLR</code>, I have to compute <code>gamma</code>.</p>
<p>The problem is that the second scheduler seems stuck at the final learning rate.</p>
<p>Here is my code:</p>
<pre><code class="lang-auto">from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR, ExponentialLR
from torch import optim
from torch import nn
import matplotlib.pyplot as plt
import numpy as np
    
max_lr = 0.01
initial_lr = max_lr / 1e4
model = nn.Linear(10,10)
opt = optim.SGD(model.parameters(), lr=initial_lr, momentum=0.9)
num_epochs = 20
gamma = np.power(max_lr/initial_lr,1/7)
warmup_scheduler = ExponentialLR(opt,
                         gamma=gamma)
cos_scheduler = CosineAnnealingLR(opt,
                                  T_max = num_epochs-8,
                                  eta_min=max_lr / 10000.)
scheduler = SequentialLR(opt,
                         schedulers=[warmup_scheduler, cos_scheduler],
                         milestones=[8])
lrs = []
for epoch in range(num_epochs):
    lr = scheduler.get_last_lr()[0]
    lrs.append(lr)
    print(f"epoch: {epoch}: lr = {lr}")
    scheduler.step()
    
fig = plt.figure(figsize=(8,4))
ax = fig.add_subplot(1,1,1)
ax.plot(np.arange(len(lrs)),lrs)
# ax.set_yscale('log')
fig.tight_layout()
plt.savefig('foo.png', bbox_inches='tight')
plt.show()
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/7/f7e4b156c44478023219f9fc657f467e59990f41.png" data-download-href="https://discuss.pytorch.org/uploads/default/f7e4b156c44478023219f9fc657f467e59990f41" title="exp"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/f/7/f7e4b156c44478023219f9fc657f467e59990f41_2_690x345.png" alt="exp" data-base62-sha1="zmXUeX3FK7nqJcgzqnsT5Tbq5mp" width="690" height="345" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/f/7/f7e4b156c44478023219f9fc657f467e59990f41_2_690x345.png, https://discuss.pytorch.org/uploads/default/original/3X/f/7/f7e4b156c44478023219f9fc657f467e59990f41.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/f/7/f7e4b156c44478023219f9fc657f467e59990f41.png 2x" data-dominant-color="FBFBFC"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">exp</span><span class="informations">800×400 14.2 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Thanks for any help.</p>