【Keras】自定义学习率更新策略和实时监控学习率变化

这里使用的是tensorflow中的keras，在import时用的是如下类似的语句：

from tensorflow.python.keras.models import Model

一、实时监控学习率

要实现实时监控学习率变化，定义这样一个函数：

def get_lr_metric(optimizer):
    def lr(y_true, y_pred):
        return optimizer.lr
    return lr

然后在model.compile的时候加入到metrics规则。

opt = SGD(lr=lr, decay=0)
lr_metric = get_lr_metric(opt)

model.compile(optimizer=opt,
              loss='categorical_crossentropy',
              metrics=['accuracy', lr_metric])

这样在训练的时候，就可以实时的观察到当前的学习率了。

二、添加回调函数

因为在优化器的实现中，并没有更新self.lr的值。

 所以如果仅仅做了上面的第一步，那得到的lr始终都是设置的初始学习率，即是设置了decay值也不会改变。

这个时候需要自定义一个回调函数。

重写Callback类。

from tensorflow.python.keras import backend as K
from tensorflow.python.keras.callbacks import Callback

class LRCallback(Callback):
    def on_epoch_end(self, epoch, logs=None):
        lr = self.model.optimizer.lr
        decay = self.model.optimizer.decay
        iterations = self.model.optimizer.iterations
        lr_with_decay = lr / (1. + decay * K.cast(iterations, K.dtype(decay)))
        return K.eval(lr_with_decay)

然后在model.fit_generator的时候加其加入到回调函数列表中。

learning_rate_callback = LRCallback()
history = model.fit_generator(
    train_generator,
    steps_per_epoch=train_generator.n//train_generator.batch_size,
    epochs=transfer_epochs,
    callbacks=[learning_rate_callback ],
    validation_data=valid_generator,
    validation_steps=valid_generator.n//valid_generator.batch_size
)

这样就可以在每个迭代看到实时更新的学习率了。

最后还可以画出整个训练过程的学习率变化曲线。

lr_range = history.history['lr']
plt.figure(figsize=(8, 8))
plt.plot(len(lr_range), lr_range, label='Learning rate')
plt.legend(loc='upper right')
plt.title('Learning rate')
plt.show()

三、自定义学习率更新策略

tf.keras.optimizers.schedules中已经有了几种学习率更新策略了。

或者像我使用的版本还没有这种策略，都不能完全满足应用的要求。所以这时候就需要自定义更新策略。

比如我想自定义cosine_decay的cycle版本。

learning_rate = 0.01
t_mul = 2.0
m_mul = 0.8
alpha = 0.00001
def scheduler_cosine_decay(epoch):
    if epoch == 0:
        global global_step
        global lr
        global sess
        global_step = tf.Variable(tf.constant(0))
        lr = tf.train.cosine_decay_restarts(learning_rate , global_step, first_decay_steps, t_mul, m_mul, alpha )
        sess = tf.Session()
    new_lr = sess.run(lr, feed_dict={global_step: epoch})
    return new_lr

函数仅有一个参数epoch。

然后把它加入回调函数列表。

from tensorflow.python.keras.callbacks import LearningRateScheduler

learning_rate_callback = LearningRateScheduler(scheduler_cosine_decay)
history = model.fit_generator(
    train_generator,
    steps_per_epoch=train_generator.n//train_generator.batch_size,
    epochs=transfer_epochs,
    callbacks=[learning_rate_callback ],
    validation_data=valid_generator,
    validation_steps=valid_generator.n//valid_generator.batch_size
)

这样就完成了自定义学习率更新策略，并且可以被实时监控。

想定义其他的自定义学习率可以参考【Tensorflow】1.x版本中的各种学习率策略，然后仿造上面的方法进行。