chemprop.schedulers

chemprop.schedulers#

Functions#

build_NoamLike_LRSched(optimizer, warmup_steps, ...)

Build a Noam-like learning rate scheduler which schedules the learning rate with a piecewise linear followed

Module Contents#

chemprop.schedulers.build_NoamLike_LRSched(optimizer, warmup_steps, cooldown_steps, init_lr, max_lr, final_lr)[source]#

Build a Noam-like learning rate scheduler which 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_steps (int) – The number of steps during which to linearly increase the learning rate.

  • cooldown_steps (int) – The number of steps during which to exponential decay the learning rate.

  • init_lr (float) – The initial learning rate.

  • max_lr (float) – The maximum learning rate (achieved after warmup_steps).

  • final_lr (float) – The final learning rate (achieved after cooldown_steps).

References