feat(train): add cosine LR scheduler with linear warmup - #13
Merged
Conversation
BoardOCR training previously ran with a constant LR (default 3e-4) for the
full epoch budget, leaving cheap image-classification gains on the table
especially for 200-epoch runs.
Add three CLI flags plumbed into a per-step SequentialLR of LinearLR +
CosineAnnealingLR:
--scheduler {none,cosine} default: cosine
--warmup-epochs INT default: 5
--min-lr-ratio FLOAT default: 0.01 (min_lr = lr * ratio)
The scheduler steps once per batch (all ranks share len(train_loader) shard
size, so DDP stays in lockstep), its state is persisted in the checkpoint,
and current LR is logged to W&B as step/lr each --log-every batch.
--scheduler=none reverts to the prior constant-LR behavior for
back-compat / A/B comparisons.
Note: version bump deferred; PR #11 already bumps to 0.3.2, so whichever
PR merges second will need to re-bump.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Companion bump for the LR scheduler feature. develop is at 0.3.2 after PR Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
tkgstrator
force-pushed
the
feat/lr-scheduler
branch
from
July 13, 2026 07:26
2413e90 to
2237d6d
Compare
Merged
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
BoardOCR training gains a per-step LR scheduler (linear warmup → cosine anneal) plus three CLI flags:
--scheduler {none,cosine}(defaultcosine)--warmup-epochs(default5)--min-lr-ratio(default0.01, somin_lr = lr * 0.01)Scheduler state persists in the checkpoint (
schedulerkey alongsideoptimizer/scaler) so resume continues on the same LR curve. The current LR is logged to W&B asstep/lreach--log-everybatch.--scheduler=nonerecovers the previous constant-LR behavior for direct A/B against existing runs.Why
Constant
3e-4for the full 200 epoch budget was leaving standard image-classification headroom on the table — cosine+warmup is the modern default for this class of model. Also useful for the new 6-GPU DDP setup (effective batch 384), where warmup smooths the first few epochs at the higher effective LR.Test plan
load_state_dictat step 10000 of a 20000-step schedule--helpsurfaces the three new flags, defaults preservedstep/lrshows the warmup + cosine curveVersion
Deferred here — PR #11 already bumps
pyproject.tomlto0.3.2. Whichever of the two PRs merges second will need to re-bump.🤖 Generated with Claude Code