DLRS: Dynamic Learning Rate Scheduler for PyTorch

DLRS is an independent implementation of the loss-driven learning rate rule in Improving neural network training using dynamic learning rate schedule for PINNs and image classification by Veerababu Dharanalakota, Ashwin Arvind Raikar, and Prasanta Kumar Ghosh. It changes an optimizer's learning rates using the first, last, and mean batch loss of an epoch.

Status: experimental / Alpha. This package is not part of PyTorch. Correct implementation of the update rule does not guarantee convergence, better accuracy, or unchanged training quality. Large-scale production readiness, broad accelerator coverage, and upstream acceptance have not been established. See the reliability audit for measured evidence and remaining gaps.

The correctness contract defines supported behavior; its verification report separates solver-checked mathematical properties, adversarial implementation tests and deliberately broken implementations. These are bounded evidence, not a proof of the entire Python/PyTorch system.

The single-GPU CUDA validation records 18 required installed-wheel checks on a free Colab T4; two native-BF16 cases were unsupported and skipped. Each run qualifies its recorded source revision only. A pinned notebook reproduces the bounded check. The readiness ledger separates completed checks from remaining deployment and upstream requirements.

An earlier revision completed 96 planned trials across MNIST, CIFAR-10, a manufactured PINN, and a small Transformer. MNIST and Transformer met their provisional comparison margins; CIFAR-10 and the PINN failed against cosine. CIFAR-10 also failed to establish the required tolerance against fixed LR. These results block a general replacement recommendation; see the workload results and Transformer results. The subsequent numerical corrections need a new declared training-quality study.

Installation

Install an appropriate PyTorch build for your platform first using the official instructions, then explicitly select this experimental prerelease:

python -m pip install --pre "pytorch-dlrs==0.3.0.dev0"

The package declares Python 3.8+ and PyTorch 2.0+; available combinations depend on PyTorch's wheels. These lower bounds are compatibility targets, not a claim that every newer release or platform has been tested. The compatibility report records executed local checks; the CI matrix and contribution guide define the ongoing checks. NumPy, torchvision, plotting tools, and test tools are optional, not runtime dependencies of dlrs.

For development from source:

git clone https://github.com/Thabhelo/pytorch-dlrs.git
cd pytorch-dlrs
python -m pip install -e '.[dev]'

Version 0.3.0.dev0 intentionally changes behavior from 0.2.1; read the migration notes before replacing an existing installation or resuming a run. For release history, consult CHANGELOG.md.

Quick start

Run scheduler.step(batch_losses) once at the end of an epoch, after optimizer updates. Losses must be finite, nonnegative scalar values in batch order. The following example is self-contained:

import torch
from dlrs import DLRSScheduler

torch.manual_seed(0)
model = torch.nn.Linear(4, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
scheduler = DLRSScheduler(optimizer, min_lr=1e-6, max_lr=0.1)
inputs = torch.randn(32, 4)
targets = torch.randn(32, 1)

for epoch in range(3):
    batch_losses = []
    for data, target in zip(inputs.split(8), targets.split(8)):
        optimizer.zero_grad()
        loss = torch.nn.functional.mse_loss(model(data), target)
        loss.backward()
        optimizer.step()
        batch_losses.append(loss.detach().item())
    scheduler.step(batch_losses)
    print(epoch, scheduler.get_last_lr())

max_lr=0.1 is illustrative; determine useful bounds with a controlled comparison on your workload. A bound limits a rate, not training risk. With mixed precision, use the unscaled loss and account for optimizer steps skipped by the scaler. With gradient accumulation, explicitly define whether each observation represents a microbatch or a completed optimizer update.

Update rule

For at least two losses, DLRS computes slope = (last - first) / mean and, independently for each positive learning rate lr:

scale = 10 ** floor(log10(lr))
delta = delta_i if slope < 0 else delta_o if slope <= 1 else delta_d
new_lr = clamp(lr - scale * delta * slope, min_lr, max_lr)

A negative slope increases the rate; a positive slope decreases it. These names describe a heuristic observation, not a proof that the model is converging or diverging. The rule uses the endpoints and mean, so batch ordering, uneven batch sizes, and noisy losses matter. Design notes describe numerical behavior and limitations.

Bounds apply during loss-driven updates; construction and skipped/no-loss steps do not rewrite initial or externally changed rates. Zero-rate groups remain zero. Invalid values are rejected before applying a partial update. Empty/single-observation and all-zero loss sequences warn and leave rates unchanged. step() without losses advances the scheduler counter but does not perform a DLRS adjustment.

Integration contract

• Use a normal eager PyTorch optimizer with scalar real learning rates. Tensor rates must be writable scalar floating-point values without gradients; their identity is preserved. Graph capture, compiled or differentiable optimizers, sharded optimizers, and arbitrary custom optimizers are not certified.
• Serialize optimizer updates, scheduler steps, and checkpoint operations for each optimizer. The scheduler provides no thread synchronization.
• Keep optimizer parameter groups fixed after constructing the scheduler. Adding, removing, or reordering groups requires a new scheduler and an explicit state migration.
• For distributed training, every rank needs identical scheduler observations, state, and resulting learning rates. DLRS performs no collectives. Calling step() only on rank 0 leaves other optimizers out of sync; see the DDP example.
• ChainedScheduler and SequentialLR do not forward batch losses. They must not be used to orchestrate DLRS. Multiple schedulers also change each other's behavior; use an explicitly tested training-loop policy if combining them.
• DLRSOnPlateau uses a rolling window of epoch metrics, not within-epoch losses. Its historical patience argument sets the window length (patience + 2), not bad-epoch patience. mode='max' reverses the slope direction. It is not a drop-in replacement for ReduceLROnPlateau.
• The Lightning example supports a deliberately narrow single-process training loop; the Transformers callback is disabled pending a correct integration. Examples are not universal framework compatibility guarantees.

Save and restore both optimizer and scheduler state. Construct the scheduler before loading either state, load the optimizer, then load the matching scheduler state. See checkpointing and migration.

Validation and examples

The local validation reports separate training modes, distributed recovery, version compatibility, vision/PINN workload quality, Transformer workload quality, ordinary-use behavior against the previous release, and scheduler scaling. A successful functional test is not evidence that an alternative learning-rate policy preserves your model quality. Use the workload acceptance results and their stated limits.

python -m pytest tests --cov=dlrs --cov-report=term-missing
ruff check dlrs tests examples benchmarks setup.py .github/scripts
mypy dlrs
python -m build
python -m twine check dist/*

To run dataset/plot examples from a checkout, install python -m pip install -e '.[examples]' with compatible torch/torchvision versions. The MNIST example, PINN example, and benchmark runner illustrate usage. They do not establish performance superiority or a safe default for an unrelated workload. Benchmark reports must record seeds, software/hardware, tuning budgets, failures, and held-out evaluation. Retain the existing scheduler as a baseline and compare quality as well as runtime before adopting DLRS.

Citation and attribution

@article{DHARANALAKOTA2025100697,
  title = {Improving neural network training using dynamic learning rate schedule for PINNs and image classification},
  author = {Dharanalakota, Veerababu and Raikar, Ashwin Arvind and Ghosh, Prasanta Kumar},
  journal = {Machine Learning with Applications},
  volume = {21},
  pages = {100697},
  year = {2025},
  doi = {10.1016/j.mlwa.2025.100697}
}

Implementation by Thabhelo. MIT License. Contribution guide, FAQ, and issues.