chemprop.schedulers

Module Contents

Classes

NoamLR

A Noam learning rate scheduler schedules the learning rate with a piecewise linear followed

class chemprop.schedulers.NoamLR(optimizer, warmup_epochs, total_epochs, steps_per_epoch, init_lrs, max_lrs, final_lrs)

Bases: torch.optim.lr_scheduler.LRScheduler

A Noam learning rate scheduler schedules the learning rate with a piecewise linear followed by an exponential decay.

The learning rate increases linearly from init_lr to max_lr over the course of the first warmup_steps then decreases exponentially to final_lr over the course of the remaining total_steps - warmup_steps (where total_steps = total_epochs * steps_per_epoch). This is roughly based on the learning rate schedule from [1], section 5.3.

Formally, the learning rate schedule is defined as:

\[\begin{split}\mathtt{lr}(i) &= \begin{cases} \mathtt{init\_lr} + \delta \cdot i &\text{if } i < \mathtt{warmup\_steps} \\ \mathtt{max\_lr} \cdot \left( \frac{\mathtt{final\_lr}}{\mathtt{max\_lr}} \right)^{\gamma(i)} &\text{otherwise} \\ \end{cases} \\ \delta &\mathrel{:=} \frac{\mathtt{max\_lr} - \mathtt{init\_lr}}{\mathtt{warmup\_steps}} \\ \gamma(i) &\mathrel{:=} \frac{i - \mathtt{warmup\_steps}}{\mathtt{total\_steps} - \mathtt{warmup\_steps}}\end{split}\]

Parameters:

• optimizer (Optimizer) – A PyTorch optimizer.

• warmup_epochs (ArrayLike) – The number of epochs during which to linearly increase the learning rate.

• total_epochs (int) – The total number of epochs.

• steps_per_epoch (int) – The number of steps (batches) per epoch.

• init_lr (ArrayLike) – The initial learning rate.

• max_lr (ArrayLike) – The maximum learning rate (achieved after warmup_epochs).

• final_lr (ArrayLike) – The final learning rate (achieved after total_epochs).

• init_lrs (numpy.typing.ArrayLike)

• max_lrs (numpy.typing.ArrayLike)

• final_lrs (numpy.typing.ArrayLike)

References

[1]

Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A.N., Kaiser, Ł. and Polosukhin, I. “Attention is all you need.” Advances in neural information processing systems, 2017, 30. https://arxiv.org/abs/1706.03762

__len__()

the number of steps in the learning rate schedule

Return type:

int

get_lr()

Get a list of the current learning rates

Return type:

numpy.ndarray

step(step=None)

Step the learning rate

Parameters:

step (int | None, default=None) – What step to set the learning rate to. If None, use self.current_step + 1.