Skip to content

Piecewise Scheduler

About

The d9d.lr_scheduler.piecewise module builds piecewise learning rate schedules, such as "warmup, hold, decay". You chain phases together instead of writing an LRScheduler subclass or a LambdaLR function for each schedule.

Phases

A schedule computes a multiplier for the base learning rate of the optimizer. It starts at initial_multiplier. Each phase moves the multiplier from the end value of the previous phase to its own target_multiplier, along a curve.

A phase can last:

  • a fixed number of steps (for_steps, or "mode": "steps" in a config);
  • until a fraction of the total steps (until_percentage, or "mode": "percentage");
  • until the end of training (fill_rest, or "mode": "rest").

The last two need total_steps. After the last phase, the multiplier stays at its end value.

Available Curves

The following curves interpolate the multiplier within a phase:

Curve Class Curve Config Description
CurveLinear "linear" Straight-line interpolation.
CurveCosine "cosine" Half-period cosine interpolation (cosine annealing).
CurvePoly(power) "poly" Polynomial interpolation. power=1 is linear, power=2 is quadratic.
CurveExponential "exponential" Exponential interpolation, linear in log space. Values below 1e-8 are treated as 1e-8.

Usage

Python API

This example builds a "linear warmup, hold, cosine decay" schedule:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
import torch
from d9d.lr_scheduler.piecewise import CurveCosine, CurveLinear, piecewise_schedule

optimizer: torch.optim.Optimizer = ...
total_steps: int = 1000

# 1. Start at 0.0.
# 2. Warm up linearly to 1.0 over 100 steps.
# 3. Hold at 1.0 until 50% of the training steps.
# 4. Decay along a cosine to 0.1 for the rest of training.
scheduler = (
    piecewise_schedule(initial_multiplier=0.0, total_steps=total_steps)
    .for_steps(100, target_multiplier=1.0, curve=CurveLinear())
    .until_percentage(0.5, target_multiplier=1.0, curve=CurveLinear())
    .fill_rest(target_multiplier=0.1, curve=CurveCosine())
    .build(optimizer)
)

Pydantic API

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
import torch
from d9d.lr_scheduler.piecewise import PiecewiseSchedulerConfig, piecewise_scheduler_from_config

optimizer: torch.optim.Optimizer = ...
total_steps: int = 1000

raw_config_json = """
{
    "initial_multiplier": 0.0,
    "phases": [
        {
            "mode": "steps",
            "steps": 100,
            "target_multiplier": 1.0,
            "curve": { "type": "linear" }
        },
        {
            "mode": "rest",
            "target_multiplier": 0.1,
            "curve": { "type": "cosine" }
        }
    ]
}
"""

scheduler_config = PiecewiseSchedulerConfig.model_validate_json(raw_config_json)

scheduler = piecewise_scheduler_from_config(
    config=scheduler_config,
    optimizer=optimizer,
    total_steps=total_steps
)

API Reference

d9d.lr_scheduler.piecewise

Piecewise learning rate schedules composed with a builder.

CurveBase

Bases: ABC

Abstract base class for interpolation curves used in scheduling.

compute(start, end, step_p) abstractmethod

Computes the interpolated value.

Parameters:

Name Type Description Default
start float

Value at the start of the phase.

 required
end float

Value at the end of the phase.

 required
step_p float

Progress through the phase, from 0.0 to 1.0.

 required

Returns:

Type Description
float

The interpolated value.

CurveCosine

Bases: CurveBase

Cosine annealing curve (half-period cosine).

CurveExponential

Bases: CurveBase

Exponential curve between the start and end values (linear in log space).

Start and end values below 1e-8 are treated as 1e-8.

CurveLinear

Bases: CurveBase

Linearly interpolates between start and end values.

CurvePoly

Bases: CurveBase

Polynomial curve along step_p ** power.

__init__(power)

Constructs the CurvePoly object.

Parameters:

Name Type Description Default
power float

Exponent of the polynomial. 1.0 is linear, 2.0 is quadratic.

 required

PiecewiseSchedulerConfig

Bases: BaseModel

Configuration for a piecewise learning rate scheduler.

Attributes:

Name Type Description
initial_multiplier float

Learning rate multiplier at step 0.
phases list[PhaseConfig]

Phase configurations, in order.

piecewise_schedule(initial_multiplier, total_steps=None)

Starts building a piecewise learning rate schedule.

Parameters:

Name Type Description Default
initial_multiplier float

Learning rate multiplier at step 0.

 required
total_steps int | None

Total number of training steps. Required for percentage-based phases.

 None

Returns:

Type Description
PiecewiseScheduleBuilder

A builder to add phases to.

piecewise_scheduler_from_config(config, optimizer, total_steps)

Builds a PyTorch learning rate scheduler from a configuration.

Parameters:

Name Type Description Default
config PiecewiseSchedulerConfig

Scheduler configuration.

 required
optimizer Optimizer

Optimizer whose learning rate the scheduler controls.

 required
total_steps int | None

Total number of training steps. Required if config has "percentage" or "rest" phases.

 required

Returns:

Type Description
LRSchedulerProtocol

The learning rate scheduler.

Raises:

Type Description
ValueError

If the phases do not fit total_steps, as described in PiecewiseScheduleBuilder, or if the type of a phase or curve configuration is unknown.