lrsched 0.1.0

Framework-agnostic learning-rate schedules as pure functions, in Python with zero dependencies.

lrsched

Framework-agnostic learning-rate schedules as pure functions, in Python with zero dependencies. Each schedule maps a step to a learning rate, so it works in any training loop or framework, or none.

Install

pip install lrsched

30-second example

from lrsched import cosine, with_warmup, sample

schedule = with_warmup(
    cosine(base_lr=1e-3, min_lr=1e-5, total_steps=1000),
    warmup_steps=100,
    start_lr=0.0,
)

lr = schedule(250)            # learning rate at step 250
curve = sample(schedule, num_steps=1000)  # the whole curve, for plotting or logging

A schedule is just Callable[[int], float]. Plug schedule(step) into your optimizer however your framework expects, or use it to drive a plain training loop.

Why this exists

Every learning-rate scheduler is tied to a framework: torch.optim.lr_scheduler, timm, transformers, or optax for JAX. If you write a custom loop, use a non-PyTorch stack, or just want to plot a schedule, you end up pasting a LambdaLR snippet. lrsched is a small, dependency-free library where each schedule is a pure function, easy to test, plot, log, and reuse anywhere.

Comparison

	lrsched	torch / timm	optax
Framework	none	PyTorch	JAX
Pure step to lr function	yes	no (optimizer-bound)	partial
Zero dependencies	yes	no	no
Composable warmup and phases	yes	partial	yes

Schedules

• constant, step_decay, multi_step, exponential
• linear, polynomial
• cosine, cosine_restarts (SGDR)
• inverse_sqrt (Transformer)
• one_cycle

Composition

• with_warmup(schedule, ...) prepends a linear warmup.
• sequential(phases) runs schedules back to back.
• sample(schedule, num_steps=...) evaluates a schedule over a range.

Parameters are required keyword arguments, so every schedule reads explicitly at the call site. Schedules hold their final value past the end rather than erroring, and a negative step raises.

Examples

python examples/schedules.py

Testing

pip install -e ".[dev]"
pytest

Tests cover the exact value of each schedule at known steps, schedule-specific shapes (restarts, the one-cycle peak, the warmup handoff), and invariants checked with Hypothesis (cosine stays within bounds, warmup is monotone).

Contributing

Issues and pull requests are welcome. See CONTRIBUTING.md.

License

MIT. See LICENSE.