如何在tensorflow/keras上实现warmup学习率策略代码示例

Warmup策略

Warmup为什么有效?
这个问题目前还没有被充分证明，我们只能从直觉上和已有的一些论文[1,2,3]得到推测：
有助于减缓模型在初始阶段对mini-batch的提前过拟合现象，保持分布的平稳
有助于保持模型深层的稳定性
可以认为，刚开始模型对数据的“分布”理解为零，或者是说“均匀分布”（当然这取决于你的初始化）；在第一轮训练的时候，每个数据点对模型来说都是新的，模型会很快地进行数据分布修正，如果这时候学习率就很大，极有可能导致开始的时候就对该数据“过拟合”，后面要通过多轮训练才能拉回来，浪费时间。当训练了一段时间（比如两轮、三轮）后，模型已经对每个数据点看过几遍了，或者说对当前的batch而言有了一些正确的先验，较大的学习率就不那么容易会使模型学偏，所以可以适当调大学习率。这个过程就可以看做是warmup。那么为什么之后还要decay呢？当模型训到一定阶段后（比如十个epoch），模型的分布就已经比较固定了，或者说能学到的新东西就比较少了。如果还沿用较大的学习率，就会破坏这种稳定性，用我们通常的话说，就是已经接近loss的local optimal了，为了靠近这个point，我们就要慢慢来。
keras提供了Callback基类(回调函数)
回调函数是一个函数的合集，会在训练的阶段中所使用。你可以使用回调函数来查看训练模型的内在状态和统计。你可以传递一个列表的回调函数（作为 callbacks 关键字参数）到 Sequential 或 Model 类型的 .fit() 方法。在训练时，相应的回调函数的方法就会被在各自的阶段被调用。
使用回调函数能够很方便的实现这样一个功能
我这里实现了一个带有warmup的指数衰减学习率,使用余弦退回也是类似的道理

使用带Warmup的指数衰减学习率策略

代码实现

from tensorflow.keras.callbacks import Callback
import tensorflow.keras.backend as K
class WarmupExponentialDecay(Callback):
    def __init__(self,lr_base=0.0002,lr_min=0.0,decay=0,warmup_epochs=0):
        self.num_passed_batchs = 0   #一个计数器
        self.warmup_epochs=warmup_epochs  
        self.lr=lr_base #learning_rate_base
        self.lr_min=lr_min #最小的起始学习率,此代码尚未实现
        self.decay=decay  #指数衰减率
        self.steps_per_epoch=0 #也是一个计数器
    def on_batch_begin(self, batch, logs=None):
        # params是模型自动传递给Callback的一些参数
        if self.steps_per_epoch==0:
            #防止跑验证集的时候呗更改了
            if self.params['steps'] == None:
                self.steps_per_epoch = np.ceil(1. * self.params['samples'] / self.params['batch_size'])
            else:
                self.steps_per_epoch = self.params['steps']
        if self.num_passed_batchs < self.steps_per_epoch * self.warmup_epochs:
            K.set_value(self.model.optimizer.lr,
                        self.lr*(self.num_passed_batchs + 1) / self.steps_per_epoch / self.warmup_epochs)
        else:
            K.set_value(self.model.optimizer.lr,
                        self.lr*((1-self.decay)**(self.num_passed_batchs-self.steps_per_epoch*self.warmup_epochs)))
        self.num_passed_batchs += 1
    def on_epoch_begin(self,epoch,logs=None):
    #用来输出学习率的,可以删除
        print("learning_rate:",K.get_value(self.model.optimizer.lr))         

使用示例:

optimizer=Adam() #学习率会被回调函数覆盖,所以这里可以不填
model=tf.keras.models.Model(inputs=inputs,outputs=outputs)
model.compile(optimizer=optimizer,loss="categorical_crossentropy",metrics=["acc"])
trained_model=model.fit(
            train_dataset,
            steps_per_epoch=12031/batch_size,
            validation_steps=20,
            validation_data=test_dataset, 
            shuffle =True,
            epochs=160,
            callbacks=[WarmupExponentialDecay(lr_base=0.0002,decay=0.00002,warmup_epochs=2)])

tensorflow版本:

def exponential_decay_with_warmup(warmup_step,learning_rate_base,global_step,learning_rate_step,learning_rate_decay,staircase=False):
'''   如果learning_rate_base=0.01或者0.1,网络可能训练失败,0.001又太小了,这时候可以考虑前10个epoch线性增长到0.1
   这样网络不会像一开始就0.1那样训练崩掉
      warmup有助于网络的收敛,在某些情况下如果没有warmup可能导致网络无法收敛
   至于理论上的解释,还没有得到充分证明,只管上来讲就是初期网络参数是随机数,对参数优化(包括BN层的方差,均值还不稳定)取决于当前这组,      
   有可能网络因此对该组数据过拟合,进入某个局部最优解,后面要多花很多组才能调回来
   前期使用小学习率,网络对数据分布大致有个了解,这时候使用大学习率就不容易跑偏
   当然训练到几十个epoch之后模型的分布就已经比较固定了，或者说能学到的新东西就比较少了。如果还沿用较大的学习率
   就会破坏这种稳定性，用我们通常的话说，就是已经接近loss的local optimal了，为了靠近这个point，我们就要慢慢来,所以后期进行
   指数衰减的学习率(或者其他的衰减形式)
   warmup_step:warmup线性增长的step数
   learning_rate_base学习率最大值
   global_step一个全局的计数器,注意设置trainable=False,并在优化器的minimize里面填写global_step=global_step
   [  global_step: tf.Variable(0,trainable=False)  ]
   learning_rate_step 每隔 step衰减一次
   learning_rate_decat 衰减率
   参数设置个人建议最终训练完学习率大致为learning_rate_base的1/50,learning_rate_decay在0.95到0.995之间
 '''
    with tf.name_scope("exponential_decay_with_warmup"):
	    linear_increase=learning_rate_base*tf.cast(global_step/warmup_step,tf.float32)
	    exponential_decay=tf.train.exponential_decay(learning_rate_base,
	                                                 global_step-warmup_step,
	                                                 learning_rate_step,
	                                                 learning_rate_decay,
	                                                 staircase=staircase)
	    learning_rate=tf.cond(global_step<=warmup_step,
	                          lambda:linear_increase,
	                          lambda:exponential_decay)
	    return learning_rate