scico.flax.train.learning_rate

Learning rate schedulers.

Functions

create_cnst_lr_schedule(config)	Create learning rate to be a constant specified value.
create_cosine_lr_schedule(config)	Create learning rate to follow a pre-specified schedule.
create_exp_lr_schedule(config)	Create learning rate schedule to have an exponential decay.

scico.flax.train.learning_rate.create_cnst_lr_schedule(config)

Create learning rate to be a constant specified value.

Parameters:

config (ConfigDict) – Dictionary of configuration. The value to use corresponds to the base_learning_rate keyword.

Return type:

Callable[[Union[Array, ndarray, bool_, number, float, int]], Union[Array, ndarray, bool_, number, float, int]]

Returns:

schedule – A function that maps step counts to values.

scico.flax.train.learning_rate.create_exp_lr_schedule(config)

Create learning rate schedule to have an exponential decay.

Parameters:

config (ConfigDict) – Dictionary of configuration. The values to use correspond to base_learning_rate, num_epochs, steps_per_epochs and lr_decay_rate.

Return type:

Callable[[Union[Array, ndarray, bool_, number, float, int]], Union[Array, ndarray, bool_, number, float, int]]

Returns:

schedule – A function that maps step counts to values.

scico.flax.train.learning_rate.create_cosine_lr_schedule(config)

Create learning rate to follow a pre-specified schedule.

Create learning rate to follow a pre-specified schedule with warmup and cosine stages.

Parameters:

config (ConfigDict) – Dictionary of configuration. The parameters to use correspond to keywords: base_learning_rate, num_epochs, warmup_epochs and steps_per_epoch.

Return type:

Callable[[Union[Array, ndarray, bool_, number, float, int]], Union[Array, ndarray, bool_, number, float, int]]

Returns:

schedule – A function that maps step counts to values.