A keyframe-style learning rate scheduler for PyTorch

username: David_Gilbertson
cooked: <p>Hello fellow PyTorch folks,</p>
<p>I’ve put together a basic LR scheduler that works quite well for my use cases. I’m sharing it here in the hope that others might find it useful too.</p>
<p>The idea is that you define a schedule with a set of frames and transitions, for example:</p>
<pre data-code-wrap="python"><code class="lang-python">scheduler = KeyframeLR(
    optimizer=my_optimizer,
    frames=[
        {"position": 0, "lr": 0.001},
        {"transition": "cos"},
        {"position": 0.2, "lr": 0.01},
        {"transition": "cos"},
        {"position": 1, "lr": 0.005},
    ],
    end=max_steps,
)
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/5/6/56be56c4a5443b72a7b9508935c94dc63f425b44.png" data-download-href="https://discuss.pytorch.org/uploads/default/56be56c4a5443b72a7b9508935c94dc63f425b44" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/5/6/56be56c4a5443b72a7b9508935c94dc63f425b44_2_690x272.png" alt="image" data-base62-sha1="cnmRqzNiMJhqSqj05Bh80moxn92" width="690" height="272" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/5/6/56be56c4a5443b72a7b9508935c94dc63f425b44_2_690x272.png, https://discuss.pytorch.org/uploads/default/optimized/3X/5/6/56be56c4a5443b72a7b9508935c94dc63f425b44_2_1035x408.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/5/6/56be56c4a5443b72a7b9508935c94dc63f425b44_2_1380x544.png 2x" data-dominant-color="26333A"><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">image</span><span class="informations">1409×557 72.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>You can have linear or cosine transitions, or pass a callable like so:</p>
<pre data-code-wrap="python"><code class="lang-python">scheduler = KeyframeLR(
    optimizer=my_optimizer,
    units="steps",
    frames=[
        {"position": 0, "lr": 0.001},
        {"position": 2_000, "lr": 0.01},
        {"transition": lambda last_lr, *_: last_lr * 0.999},
    ],
    end=10_000,
)
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/b/bbbfb0eb9c063ebb638777d3cbd5958e06270493.png" data-download-href="https://discuss.pytorch.org/uploads/default/bbbfb0eb9c063ebb638777d3cbd5958e06270493" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/b/bbbfb0eb9c063ebb638777d3cbd5958e06270493_2_690x288.png" alt="image" data-base62-sha1="qMTYXChjeXnq8OYwIbya2CMOtZF" width="690" height="288" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/b/bbbfb0eb9c063ebb638777d3cbd5958e06270493_2_690x288.png, https://discuss.pytorch.org/uploads/default/optimized/3X/b/b/bbbfb0eb9c063ebb638777d3cbd5958e06270493_2_1035x432.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/b/b/bbbfb0eb9c063ebb638777d3cbd5958e06270493_2_1380x576.png 2x" data-dominant-color="26333A"><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">image</span><span class="informations">1406×587 69.6 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>Units can be percentages, steps or even time.</p>
<p>The <a href="https://gist.github.com/davidgilbertson/2a6ac54ad6629a37e8f4d0539f7ef7bc" rel="noopener nofollow ugc">code is here</a> if you’d like to take it for a spin.</p>
<p>At the moment it’s just a semi-polished experiment/personal utility so there isn’t any documentation, but I’ve written <a href="https://david-gilbertson.medium.com/a-keyframe-style-learning-rate-scheduler-for-pytorch-b889110dcde8?sk=b3aee640e3c44a2c6ab647aa4364f626" rel="noopener nofollow ugc">a short article that explains most aspects of it</a>.</p>
<p>Feedback is most welcome!</p>

username: ptrblck
cooked: <p>Thanks for sharing your repository!</p>