Schedules

Learning rate schedules.

nabla.nn.optim.schedules.constant_schedule(initial_lr=0.001)

Constant learning rate schedule.

Parameters:

initial_lr (float) – The learning rate to maintain

Returns:

Function that takes epoch and returns learning rate

Return type:

Callable[[int], float]

nabla.nn.optim.schedules.exponential_decay_schedule(initial_lr=0.001, decay_factor=0.95, decay_every=1000)

Exponential decay learning rate schedule.

Parameters:

• initial_lr (float) – Initial learning rate

• decay_factor (float) – Factor to multiply learning rate by

• decay_every (int) – Apply decay every N epochs

Returns:

Function that takes epoch and returns learning rate

Return type:

Callable[[int], float]

nabla.nn.optim.schedules.step_decay_schedule(initial_lr=0.001, decay_factor=0.1, step_size=30)

Step decay learning rate schedule.

Parameters:

• initial_lr (float) – Initial learning rate

• decay_factor (float) – Factor to multiply learning rate by at each step

• step_size (int) – Number of epochs between each decay step

Returns:

Function that takes epoch and returns learning rate

Return type:

Callable[[int], float]

nabla.nn.optim.schedules.cosine_annealing_schedule(initial_lr=0.001, min_lr=1e-06, period=1000)

Cosine annealing learning rate schedule.

Parameters:

• initial_lr (float) – Initial learning rate

• min_lr (float) – Minimum learning rate

• period (int) – Number of epochs for one complete cosine cycle

Returns:

Function that takes epoch and returns learning rate

Return type:

Callable[[int], float]

nabla.nn.optim.schedules.warmup_cosine_schedule(initial_lr=0.001, warmup_epochs=100, total_epochs=1000, min_lr=1e-06)

Warmup followed by cosine annealing schedule.

Parameters:

• initial_lr (float) – Peak learning rate after warmup

• warmup_epochs (int) – Number of epochs for linear warmup

• total_epochs (int) – Total number of training epochs

• min_lr (float) – Minimum learning rate

Returns:

Function that takes epoch and returns learning rate

Return type:

Callable[[int], float]

nabla.nn.optim.schedules.learning_rate_schedule(epoch, initial_lr=0.001, decay_factor=0.95, decay_every=1000)

Learning rate schedule for complex function learning.

This is the original function from mlp_train_jit.py for backward compatibility. Consider using exponential_decay_schedule instead for new code.

Parameters:

• epoch (int) – Current epoch number

• initial_lr (float) – Initial learning rate

• decay_factor (float) – Factor to multiply learning rate by

• decay_every (int) – Apply decay every N epochs

Returns:

Learning rate for the current epoch

Return type:

float