Skip to content

Adam Optimizer Config

popinn.AdamConfig ¤

Configuration for an Adam optimization phase.

Source code in src/popinn/train.py 
15
16
17
18
19
20
21
22
23
24
@dataclasses.dataclass(frozen=True)
class AdamConfig:
    """Configuration for an Adam optimization phase."""

    lr: float | optax.Schedule = 1e-3
    num_epochs: int = 1000
    log_every: int = 500
    b1: float = 0.9
    b2: float = 0.999
    eps: float = 1e-8
__init__(lr=0.001, num_epochs=1000, log_every=500, b1=0.9, b2=0.999, eps=1e-08) ¤

Optimizer is optax.adam.

Parameters:

Name Type Description Default
lr float | optax.Schedule

Constant learning rate, or an optax schedule mapping step count to learning rate. Passed straight to optax.adam. When using a step-based schedule (e.g. cosine decay), its decay horizon should generally match num_epochs.

 0.001
num_epochs int

Number of optimization steps in this phase.

 1000
log_every int

Print a log line every this many steps.

 500
b1 float

Adam's first-moment decay rate.

 0.9
b2 float

Adam's second-moment decay rate.

 0.999
eps float

Adam's epsilon for numerical stability.

 1e-08

popinn.warmup_cosine(peak_lr, num_epochs, init_lr=1e-05, warmup_steps=500) ¤

Convenience wrapper for optax.warmup_cosine_decay_schedule.

Parameters:

Name Type Description Default
peak_lr float

peak learning rate reached at the end of warmup.

 required
num_epochs int

total steps; the cosine decay spans this horizon.

 required
init_lr float

learning rate at step 0.

 1e-05
warmup_steps int

number of steps spent ramping up to peak_lr.

 500

Returns:

Type Description
optax.Schedule

a callable step -> learning rate.

Raises:

Type Description
ValueError

if warmup_steps >= num_epochs (no steps left to decay over). optax's decay span is num_epochs - warmup_steps, which must be positive.
Source code in src/popinn/train.py 
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
def warmup_cosine(
    peak_lr: float,
    num_epochs: int,
    init_lr: float = 1e-5,
    warmup_steps: int = 500,
) -> optax.Schedule:
    """Convenience wrapper for `optax.warmup_cosine_decay_schedule`.

    Args:
        peak_lr (float): peak learning rate reached at the end of warmup.
        num_epochs (int): total steps; the cosine decay spans this horizon.
        init_lr (float): learning rate at step 0.
        warmup_steps (int): number of steps spent ramping up to peak_lr.

    Returns:
        (optax.Schedule): a callable step -> learning rate.

    Raises:
        ValueError: if warmup_steps >= num_epochs (no steps left to decay
            over). optax's decay span is num_epochs - warmup_steps, which
            must be positive.
    """
    if warmup_steps >= num_epochs:
        raise ValueError(f"warmup_steps ({warmup_steps}) must be < num_epochs ({num_epochs}); the cosine decay spans the remaining steps.")
    return optax.warmup_cosine_decay_schedule(
        init_value=init_lr,
        peak_value=peak_lr,
        warmup_steps=warmup_steps,
        decay_steps=num_epochs,
    )