Creates an optimizer with learning rate schedule.
(
init_lr,
num_train_steps,
num_warmup_steps,
min_lr_ratio=0.0,
adam_epsilon=1e-8,
weight_decay_rate=0.0,
include_in_weight_decay=None,
)
| 59 | |
| 60 | |
| 61 | def create_optimizer( |
| 62 | init_lr, |
| 63 | num_train_steps, |
| 64 | num_warmup_steps, |
| 65 | min_lr_ratio=0.0, |
| 66 | adam_epsilon=1e-8, |
| 67 | weight_decay_rate=0.0, |
| 68 | include_in_weight_decay=None, |
| 69 | ): |
| 70 | """Creates an optimizer with learning rate schedule.""" |
| 71 | # Implements linear decay of the learning rate. |
| 72 | lr_schedule = tf.keras.optimizers.schedules.PolynomialDecay( |
| 73 | initial_learning_rate=init_lr, |
| 74 | decay_steps=num_train_steps - num_warmup_steps, |
| 75 | end_learning_rate=init_lr * min_lr_ratio, |
| 76 | ) |
| 77 | if num_warmup_steps: |
| 78 | lr_schedule = WarmUp( |
| 79 | initial_learning_rate=init_lr, decay_schedule_fn=lr_schedule, warmup_steps=num_warmup_steps, |
| 80 | ) |
| 81 | if weight_decay_rate > 0.0: |
| 82 | optimizer = AdamWeightDecay( |
| 83 | learning_rate=lr_schedule, |
| 84 | weight_decay_rate=weight_decay_rate, |
| 85 | beta_1=0.9, |
| 86 | beta_2=0.999, |
| 87 | epsilon=adam_epsilon, |
| 88 | exclude_from_weight_decay=["LayerNorm", "layer_norm", "bias"], |
| 89 | include_in_weight_decay=include_in_weight_decay, |
| 90 | ) |
| 91 | else: |
| 92 | optimizer = tf.keras.optimizers.Adam(learning_rate=lr_schedule, epsilon=adam_epsilon) |
| 93 | # We return the optimizer and the LR scheduler in order to better track the |
| 94 | # evolution of the LR independently of the optimizer. |
| 95 | return optimizer, lr_schedule |
| 96 | |
| 97 | |
| 98 | class AdamWeightDecay(tf.keras.optimizers.Adam): |