# Optimizer and Learning Rate Scheduler This page describes the [optimization package](https://github.com/tensorflow/models/tree/master/official/modeling/optimization/) for Tensorflow Official Models (TFM) which includes optimizers, and learning rate schedulers. ## Building Optimizer and LR Scheduler We use an Optimizer factory class to manage optimizer and learning rate creation. Optimizer factory takes a config as an input, and it has member functions that are used to build optimizer and learning rate schedule. To create an optimizer and a LR schedule through OptimizerFactory, you need to do the following: 1. Define optimization config, this includes optimizer, and learning rate schedule. 2. Initialize the OptimizerFactory instance using the optimization config. 3. Build the learning rate, and the optimizer using the class member functions. The following is an example for creating an SGD optimizer with stepwise LR scheduler with linear warmup: ```python from third_party.tensorflow_models.official.modeling import optimization params = {'optimizer': { 'type': 'sgd', 'sgd': {'momentum': 0.9}}, 'learning_rate': {'type': 'stepwise', 'stepwise': { 'boundaries': [10000, 20000], 'values': [0.1, 0.01, 0.001]}}, 'warmup': {'type': 'linear', 'linear': {'warmup_steps': 500, 'warmup_learning_rate': 0.01}}} # Defines optimization config from a dictionary. opt_config = optimization.OptimizationConfig(params) # Initializes an optimization factory from optimization config. opt_factory = optimization.OptimizerFactory(opt_config) # Builds the desired learning rate scheduling instance. lr = opt_factory.build_learning_rate() # Builds the optimizer instance with the desired learning rate schedule. optimizer = opt_factory.build_optimizer(lr) ``` To initialize an OptimizerFactory, `optimizer` and `learning_rate` fields must be defined, while `warmup` is an optional field. The field `type` is used to define the type of each optimization component. The set of available types is explained in details in the following sections. In the following sections, we explain how to create different optimizers, learning rate, and warmup schedulers. We also explain how to add new optimizers, or learning rate schedulers. ## Optimizers The list of supported optimizers can be found [here](https://github.com/tensorflow/models/blob/7f239d8ec19b5c2d44e0d5aa2a09dbea0da6d737/official/modeling/optimization/optimizer_factory.py#L43). ```python OPTIMIZERS_CLS = { 'sgd': tf.keras.optimizers.SGD, 'adam': tf.keras.optimizers.Adam, 'adamw': nlp_optimization.AdamWeightDecay, 'lamb': tfa_optimizers.LAMB, 'rmsprop': tf.keras.optimizers.RMSprop } ``` You can specify the type of optimizer to be one of the above using [oneof](https://github.com/tensorflow/models/blob/master/official/modeling/hyperparams/oneof.py) config. The available config fields can be found [here](https://github.com/tensorflow/models/blob/master/official/modeling/optimization/configs/optimizer_config.py). All optimizers support gradient clipping methods: clip by value, clip by norm, clip by global norm. To speicify which method to use, you need to specify the appropiate field list [here](https://github.com/tensorflow/models/blob/7f239d8ec19b5c2d44e0d5aa2a09dbea0da6d737/official/modeling/optimization/configs/optimizer_config.py#L38). ### Example We will specify an rmsprop optimizer with discounting factor (rho) of 0.9, and global norm gradient clipping of 10.0. Below is the config to be used. ```python params = {'optimizer': { 'type': 'rmsprop', 'rmsprop': {'rho': 0.9, 'global_clipnorm': 10.0}}} ``` ### Adding a New Optimizer To add a new optimizer, you need to do the following: 1. Create a [custom](https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Optimizer#creating_a_custom_optimizer_2) of tf.keras.optimizers.Optimizer. 2. Add the required config fields under [optimization/configs/optimizer_config.py](https://github.com/tensorflow/models/blob/master/official/modeling/optimization/configs/optimizer_config.py). 3. Add the optimizer class to the list of available optimizer classes in [optimizer_factor](https://github.com/tensorflow/models/blob/master/official/modeling/optimization/optimizer_factory.py) ## Learning Rate and Warmup Schedules Learning rate with an optional warmup can be configured by specifying `learning_rate`, and `warmup` fields in optimization config. `learning_rate` is a required field, while `warmup` is an optional one. The list of supported `learning_rate` and `warmup` schedules can be found [here](https://github.com/tensorflow/models/blob/7f239d8ec19b5c2d44e0d5aa2a09dbea0da6d737/official/modeling/optimization/optimizer_factory.py#L51). ```python LR_CLS = { 'stepwise': tf.keras.optimizers.schedules.PiecewiseConstantDecay, 'polynomial': tf.keras.optimizers.schedules.PolynomialDecay, 'exponential': tf.keras.optimizers.schedules.ExponentialDecay, 'cosine': tf.keras.experimental.CosineDecay, 'power': lr_schedule.DirectPowerDecay, } WARMUP_CLS = { 'linear': lr_schedule.LinearWarmup, 'polynomial': lr_schedule.PolynomialWarmUp } ``` In addition, a `constant` learning rate can be specified. ## How Learning Rate Works Learning rate takes `step` as an input, and it returns the learning rate value. As the training progresses, usually learning rate value decays. Warmup schedule is often used to stablize the training. Warmup schedule starts from a low learning rate value, and it gradually increases until it reaches the initial value for the regular learning rate decay schedule. We combine `learning_rate` (lr) with `warmup` (warmup) schedules as follows * Steps [0, warmup_steps): `learning_rate = warmup(step)` * Steps [warmup_steps, train_steps): `learning_rate = lr(step)` * We designed the warmup schedule such that final warmup learning rate is inferred from the learning rate schedule (i.e. `learning_rate(warmup_steps) = warmup(warmup_steps)`). Note that, warmup schedule doesn't delay the regular learning rate decay by warmup_steps, instead it replaces it. Learning rate value is logged every [summary_interval](https://github.com/tensorflow/models/blob/7f239d8ec19b5c2d44e0d5aa2a09dbea0da6d737/official/core/config_definitions.py#L181). If warmup_steps are less that the `summary_interval`, you won't be able to see warmup values. ### Example We want to specify a cosine learning rate decay with decay_steps of 20000, with a linear warmup schedule for the first 500 steps. ```python params = {'learning_rate': {'type': 'cosine', 'cosine': {'decay_steps': 20000}}, 'warmup': {'type': 'linear', 'linear': {'warmup_steps': 500}}} ``` ## Customizing Optimizer Inside Task Optimizer and learning rate are created inside the [task](https://github.com/tensorflow/models/blob/7f239d8ec19b5c2d44e0d5aa2a09dbea0da6d737/official/core/base_task.py#L73). If different optimizers/learning rate schedulers are needed, they can be defined by overriding the class method. ## Important Factors To Consider * Batch size: Changing batch size usually requires scaling learning rate values, and number of training steps. Make sure that you change appropriate values as batch size changes. * Train steps: Train steps is highly correlated with fields such as `decay_steps` for cosine learning rate decay. Changing one without changing the other might result in undesired behavior.