spacetimeformer.lr_scheduler package

class spacetimeformer.lr_scheduler.lr_scheduler.LearningRateScheduler(optimizer, lr)

Bases: torch.optim.lr_scheduler._LRScheduler

Provides inteface of learning rate scheduler.

Note

Do not use this class directly, use one of the sub classes.

get_lr()

static set_lr(optimizer, lr)

step(*args, **kwargs)

class spacetimeformer.lr_scheduler.reduce_lr_on_plateau_lr_scheduler.ReduceLROnPlateauScheduler(optimizer: torch.optim.optimizer.Optimizer, lr: float, patience: int = 1, factor: float = 0.3)

Bases: spacetimeformer.lr_scheduler.lr_scheduler.LearningRateScheduler

Reduce learning rate when a metric has stopped improving. Models often benefit from reducing the learning rate by a factor of 2-10 once learning stagnates. This scheduler reads a metrics quantity and if no improvement is seen for a ‘patience’ number of epochs, the learning rate is reduced.

Parameters

• optimizer (Optimizer) – Optimizer.

• lr (float) – Initial learning rate.

• patience (int) – Number of epochs with no improvement after which learning rate will be reduced.

• factor (float) – Factor by which the learning rate will be reduced. new_lr = lr * factor.

step(val_loss: float)

class spacetimeformer.lr_scheduler.transformer_lr_scheduler.TransformerLRScheduler(optimizer: torch.optim.optimizer.Optimizer, init_lr: float, peak_lr: float, final_lr: float, final_lr_scale: float, warmup_steps: int, decay_steps: int)

Bases: spacetimeformer.lr_scheduler.lr_scheduler.LearningRateScheduler

Transformer Learning Rate Scheduler proposed in “Attention Is All You Need”

Parameters

• optimizer (Optimizer) – Optimizer.

• init_lr (float) – Initial learning rate.

• peak_lr (float) – Maximum learning rate.

• final_lr (float) – Final learning rate.

• final_lr_scale (float) – Final learning rate scale

• warmup_steps (int) – Warmup the learning rate linearly for the first N updates

• decay_steps (int) – Steps in decay stages

step(val_loss: Optional[torch.FloatTensor] = None)

class spacetimeformer.lr_scheduler.warmup_lr_scheduler.WarmupLRScheduler(optimizer: torch.optim.optimizer.Optimizer, init_lr: float, peak_lr: float, warmup_steps: int)

Bases: spacetimeformer.lr_scheduler.lr_scheduler.LearningRateScheduler

Warmup learning rate until total_steps

Parameters

optimizer (Optimizer) – wrapped optimizer.

step(val_loss: Optional[torch.FloatTensor] = None)

class spacetimeformer.lr_scheduler.warmup_reduce_lr_on_plateau_scheduler.WarmupReduceLROnPlateauScheduler(optimizer: torch.optim.optimizer.Optimizer, init_lr: float, peak_lr: float, warmup_steps: int, patience: int = 1, factor: float = 0.3)

Bases: spacetimeformer.lr_scheduler.lr_scheduler.LearningRateScheduler, torch.optim.lr_scheduler.ReduceLROnPlateau

Warmup learning rate until warmup_steps and reduce learning rate on plateau after.

Parameters

• optimizer (Optimizer) – wrapped optimizer.

• init_lr (float) – Initial learning rate.

• peak_lr (float) – Maximum learning rate.

• warmup_steps (int) – Warmup the learning rate linearly for the first N updates.

• patience (int) – Number of epochs with no improvement after which learning rate will be reduced.

• factor (float) – Factor by which the learning rate will be reduced. new_lr = lr * factor.

load_state_dict(state_dict)

Loads the schedulers state.

Parameters

state_dict (dict) – scheduler state. Should be an object returned from a call to state_dict().

state_dict()

Returns the state of the scheduler as a dict.

It contains an entry for every variable in self.__dict__ which is not the optimizer.

step(val_loss: Optional[float] = None)