-
Notifications
You must be signed in to change notification settings - Fork 1
/
optimization.py
42 lines (35 loc) · 2.4 KB
/
optimization.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
import tensorflow as tf
class WarmUpThenLinDecaySchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
"""
Learning rate schedule that first increases the learning rate linearly from 0 to initial_learning_rate
and then decreases it linearly to zero. Note that there is a step when the schedule switches, as the
linear decrease starts at step 0 and is therefore smaller than initial_learning_rate after the warmup steps.
This is done to have the schedule equivalent to the one used in the original BERT code.
See https://github.com/google-research/bert/blob/master/optimization.py for reference. (The step in learning
rate is not explicitly mentioned there, but a closer inspection of the code reveals it).
"""
def __init__(self, initial_learning_rate, warm_up_steps, total_steps):
self.warm_up_steps = warm_up_steps
self.warm_up_schedule = tf.keras.optimizers.schedules.PolynomialDecay(initial_learning_rate=0.0,
decay_steps=warm_up_steps,
end_learning_rate=initial_learning_rate)
self.decay_schedule = tf.keras.optimizers.schedules.PolynomialDecay(initial_learning_rate=initial_learning_rate,
decay_steps=total_steps,
end_learning_rate=0.0)
def __call__(self, step):
return tf.cond(step < self.warm_up_steps,
lambda: self.warm_up_schedule(step), lambda: self.decay_schedule(step))
GRAD_CLIP_NORM = 1.0 # default in BERT, see https://github.com/google-research/bert/blob/master/optimization.py
class ClipAdam(tf.keras.optimizers.Adam):
"""
Small workaround wrapper to clip gradients, as gradient clipping is otherwise ignored in Keras TF2.1 and bug
fix will only ship with TF2.2.
See https://github.com/tensorflow/tensorflow/commit/69da929ad4d5ba605507efa1f52b382a55b6a969 for reference.
"""
def apply_gradients(self, grads_and_vars, name=None, **kwargs):
grads = []
vars = []
for grad, var in grads_and_vars:
grads.append(tf.clip_by_norm(grad, GRAD_CLIP_NORM))
vars.append(var)
return super(ClipAdam, self).apply_gradients(zip(grads, vars), name=name, **kwargs)