LinearLR

LinearLR#

class braintools.optim.LinearLR(base_lr=0.001, start_factor=0.3333333333333333, end_factor=1.0, total_iters=5, last_epoch=0)#

Linear learning rate scheduler - Linearly scales learning rate between two factors.

LinearLR multiplies the base learning rate by a factor that changes linearly from start_factor to end_factor over total_iters epochs. This is commonly used for learning rate warmup or cooldown phases in training.

Parameters:

start_factor (float) – Multiplicative factor for the learning rate at the start (epoch 0). The initial lr will be base_lr * start_factor. Must be in range (0, 1]. Default: 1/3.
end_factor (float) – Multiplicative factor for the learning rate at the end (after total_iters). The final lr will be base_lr * end_factor. Must be in range (0, 1]. Default: 1.0.
total_iters (int) – Number of epochs over which to linearly transition from start_factor to end_factor. Default: 5.
last_epoch (int) – The index of the last epoch. Used for resuming training. Default: -1 (starts from beginning).

Notes

The learning rate at epoch \(t\) is computed as:

\[\begin{split}\eta_t = \begin{cases} \eta_0 \cdot s & \text{if } t = 0 \\ \eta_0 \cdot e & \text{if } t > T \\ \eta_0 \cdot \left(s + (e - s) \cdot \frac{t}{T}\right) & \text{otherwise} \end{cases}\end{split}\]

where \(\eta_0\) is the base learning rate, \(s\) is start_factor, \(e\) is end_factor, \(T\) is total_iters, and \(t\) is the current epoch.

Key characteristics:

Smooth linear transition between two learning rate values
Most commonly used for warmup (start_factor < end_factor)
Can also be used for cooldown (start_factor > end_factor)
Simple and predictable learning rate schedule

Common usage patterns:

Warmup: start_factor=0.01, end_factor=1.0, total_iters=5-10
Cooldown: start_factor=1.0, end_factor=0.1, total_iters=10-20
Gradual increase: start_factor=0.1, end_factor=1.0, total_iters=100

Examples

Learning rate warmup:

>>> import braintools
>>> import brainstate
>>>
>>> model = brainstate.nn.Linear(10, 5)
>>> # Warmup from 0.001 * 0.1 = 0.0001 to 0.001 over 10 epochs
>>> scheduler = braintools.optim.LinearLR(
...     start_factor=0.1,
...     end_factor=1.0,
...     total_iters=10
... )
>>> optimizer = braintools.optim.Adam(lr=scheduler)
>>> optimizer.register_trainable_weights(model.states(brainstate.ParamState))
>>>
>>> for epoch in range(15):
...     # Training code
...     scheduler.step()
...     if epoch < 11:
...         print(f"Epoch {epoch}: lr ≈ {optimizer.current_lr:.6f}")
# lr gradually increases from 0.0001 to 0.001

Standard warmup with default parameters:

>>> # Default: warmup from base_lr/3 to base_lr over 5 epochs
>>> scheduler = braintools.optim.LinearLR()
>>> optimizer = braintools.optim.SGD(lr=scheduler, momentum=0.9)
>>> optimizer.register_trainable_weights(model.states(brainstate.ParamState))
>>>
>>> # lr increases from ~0.00033 to 0.001 over 5 epochs

Learning rate cooldown:

>>> # Linearly decrease lr from base_lr to base_lr*0.01 over 20 epochs
>>> scheduler = braintools.optim.LinearLR(
...     start_factor=1.0,
...     end_factor=0.01,
...     total_iters=20
... )
>>> optimizer = braintools.optim.Adam(lr=scheduler)
>>> optimizer.register_trainable_weights(model.states(brainstate.ParamState))
>>>
>>> for epoch in range(30):
...     optimizer.step(grads)
...     scheduler.step()
# lr decreases from 0.001 to 0.00001 over first 20 epochs, then stays at 0.00001

Combining with StepLR for warmup + decay:

>>> # Warmup for 5 epochs, then step decay
>>> warmup = braintools.optim.LinearLR(
...     start_factor=0.1,
...     end_factor=1.0,
...     total_iters=5
... )
>>> decay = braintools.optim.StepLR(base_lr=0.01, step_size=30, gamma=0.1)
>>> scheduler = braintools.optim.ChainedScheduler([warmup, decay])
>>>
>>> optimizer = braintools.optim.SGD(lr=scheduler, momentum=0.9)
>>> optimizer.register_trainable_weights(model.states(brainstate.ParamState))
>>>
>>> for epoch in range(90):
...     optimizer.step(grads)
...     scheduler.step()

Gradual learning rate increase:

>>> # Start with very small lr and gradually increase
>>> scheduler = braintools.optim.LinearLR(
...     start_factor=0.01,
...     end_factor=1.0,
...     total_iters=100
... )
>>> optimizer = braintools.optim.Adam(lr=scheduler)
>>> optimizer.register_trainable_weights(model.states(brainstate.ParamState))
>>>
>>> # lr increases from 0.00001 to 0.001 over 100 epochs

Fine-tuning with gentle start:

>>> # Start at 30% of base lr, reach full lr in 3 epochs
>>> scheduler = braintools.optim.LinearLR(
...     start_factor=0.3,
...     end_factor=1.0,
...     total_iters=3
... )
>>> optimizer = braintools.optim.Adam(lr=scheduler)
>>> optimizer.register_trainable_weights(model.states(brainstate.ParamState))
>>>
>>> for epoch in range(20):
...     finetune_epoch(model, optimizer, finetune_loader)
...     scheduler.step()

LinearLR

Contents

LinearLR#