基于keras的transformer learning rate schedule

大家都知道，keras的learning rate schedule是基于epoch的，对于基于steps的learning rate schedule来说，比较难实现，网上都是实现了的tf2的版本的，对于tf1版本的几乎没有，因此我写了一个基于keras2.3.1以及tf1.15的transformer learning rate schedule

        for p, g, m, v, vhat in zip(params, grads, ms, vs, vhats):
            m_t = (self.beta_1 * m) + (1. - self.beta_1) * g
            v_t = (self.beta_2 * v) + (1. - self.beta_2) * K.square(g)
            if self.amsgrad:
                vhat_t = K.maximum(vhat, v_t)
                p_t = p - lr_t * m_t / (K.sqrt(vhat_t) + self.epsilon)
                self.updates.append(K.update(vhat, vhat_t))
            else:
                p_t = p - lr_t * m_t / (K.sqrt(v_t) + self.epsilon)

            self.updates.append(K.update(m, m_t))
            self.updates.append(K.update(v, v_t))
            new_p = p_t

            # Apply constraints.
            if getattr(p, 'constraint', None) is not None:
                new_p = p.constraint(new_p)

            self.updates.append(K.update(p, new_p))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

在keras.optimizer.Adam中，我们可以看到，p_t = p - lr_t * m_t / (K.sqrt(v_t) + self.epsilon)，因此我们要做的就是改变lr_t。lr_t是输入Adam的learning_rate，没办法轻易改变，因此我的想法是将他固定为1，再在每次更新的时候，重新乘以一个新的lr_multiplier，这个lr_multiplier即为transformer learning rate。

		@K.symbolic
        def get_updates(self, loss, params):
            lr_multiplier = transformer_schedule(self.iterations,
                                                      self.start_step,
                                                      self.warmup_steps,
                                                      self.d_model)

            old_update = K.update

            def new_update(x, new_x):
                if is_one_of(x, params):
                    new_x = x + (new_x - x) * lr_multiplier
                return old_update(x, new_x)

            K.update = new_update
            updates = super(NewOptimizer, self).get_updates(loss, params)
            K.update = old_update

            return updates
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

我们的主要做法是通过设定一个新的new_update函数，将p_t = p - lr_t * m_t / (K.sqrt(v_t) + self.epsilon)改写成p + (p_t-p)*lr_multiplier = p - lr_t * lr_multiplier * m_t / (K.sqrt(v_t) + self.epsilon)。如此一来，keras中的transformer learning rate schedule就成功实现了。

完整版本的keras transformer learning rate schedule已经开源在了keras-transformer-schedual，欢迎大家使用