keras自定义优化器

自定义一个SGD优化器

from keras.legacy import interfaces
from keras.optimizers import Optimizer
from keras import backend as K
class SGD(Optimizer):
    def __init__(self,lr=0.01,**kwargs):
        super(SGD,self).__init__(**kwargs)
        with K.name_scope(self.__name__):
            self.iterations = K.variable(0, dtype='int64', name='iterations')
            self.lr = K.variable(lr, name='lr')
            
    @interfaces.legacy_get_updates_support
    def get_updates(self, loss, params):
        grads = self.get_gradients(loss, params) #获取梯度
        self.updates = [K.update_add(self.iterations, 1)] # 定义赋值算子集合
        self.weights = [self.iterations] #优化器带来的权重，在保存模型时会被保存
        for p, g in zip(params, grads):
            new_p =p -self.lr*g
            #如果有约束，对参数加上约束
            if getattr(p, 'constraint', None) is not None: 
                new_p = p.constraint(new_p)
            #添加赋值
            self.updates.append(K.update(p, new_p))
        return self.updates
    def get_config(self):
        config = {'lr': float(K.get_value(self.lr))}
        base_config = super(SGD,self).get_config()
        return dict(list(base_config.items()) + list(config.items()))
 

实现“软batch”

假如模型比较庞大，自己的显卡最多也就能跑 batch size=16，但又想起到 batch size=64 的效果，那可以怎么办呢？
每次算 batch size=16，然后把梯度缓存起来，4 个 batch 后才更新参数。也就是说，每个小batch 都算梯度，但每 4 个 batch 才更新一次参数。

class MySGD(Optimizer):
    """
    Keras中简单自定义SGD优化器每隔一定的batch才更新一次参数
    """
    def __init__(self, lr=0.01, steps_per_update=1, **kwargs):
        super(MySGD, self).__init__(**kwargs)
        
        with K.name_scope(self.__class__.__name__):
            self.iterations = K.variable(0, dtype='int64', name='iterations')
            self.lr = K.variable(lr, name='lr')
            self.steps_per_update = steps_per_update #多少batch才更新一次
            
    @interfaces.legacy_get_updates_support
    def get_updates(self, loss, params):
        """
        主要的参数更新算法
        """
        shapes = [K.int_shape(p) for p in params]
        sum_grads = [K.zeros(shape) for shape in shapes] # 平均梯度，用来梯度下降
        grads = self.get_gradients(loss, params) # 当前batch梯度
        self.updates = [K.update_add(self.iterations, 1)] # 定义赋值算子集合
        self.weights = [self.iterations] + sum_grads # 优化器带来的权重，在保存模型时会被保存
        for p, g, sg in zip(params,grads,sum_grads):
            #梯度下降
            new_p = p - self.lr * sg / float(self.steps_per_update)
            
            #如果有约束，对参数加上约束
            if getattr(p, 'constraint', None) is not None:
                new_p = p.constraint(new_p)
            cond = K.equal(self.iterations % self.steps_per_update, 0)
            
            #满足条件才更新参数
            self.updates.append(K.switch(cond, K.update(p, new_p), p))
            
            #满足条件就要重新累积，不满足条件直接累积
            self.updates.append(K.switch(cond, K.update(sg, g), K.update(sg, sg+g)))
            
        return self.updates
    def get_config(self):
        config = {'lr': float(K.get_value(self.lr)),'steps_per_update': self.steps_per_update}
        base_config = super(MySGD, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

“侵入式”优化器

image.png

其中 ,p 是参数向量，g 是梯度， 表示 p 的第 i 次迭代时的结果。
这个算法需要走两步，大概意思就是普通的梯度下降先走一步（探路），然后根据探路的结果取平均，得到更精准的步伐，等价地可以改写为：

image.png

但是实现这类算法却有个难题，要计算两次梯度，一次对参数 ，另一次对参数。而前面的优化器定义中 get_updates 这个方法却只能执行一步（对应到 tf 框架中，就是执行一步 sess.run，熟悉 tf 的朋友知道单单执行一步 sess.run 很难实现这个需求），因此实现不了这种算法。

class HeunOptimizer:
    """
    自定义Keras的侵入式优化器
    """
    def __init__(self, lr):
        self.lr = lr
    def __call__(self, model):
        """
        需要传入模型，直接修改模型的训练函数，而不按常规流程使用优化器，所以称为“侵入式”
        下面的大部分代码，都是直接抄自keras的源码：
        https://github.com/keras-team/keras/blob/master/keras/engine/training.py#L491
        也就是keras中的_make_train_function函数
        """
        params = model._collected_trainable_weights
        loss = model.total_loss
        inputs = (model._feed_inputs + model._feed_targets + model._feed_sample_weights)
        inputs += [K.learning_phase()]
        with K.name_scope('training'):
            
            with K.name_scope('heun_optimizer'):
                
                old_grads = [[K.zeros(K.int_shape(p)) for p in params]]
                update_functions = []
                
                for i,step in enumerate([self.step1, self.step2]):
                    updates = (model.updates + step(loss, params, old_grads) + model.metrics_updates)
                    #给每一步定义一个K.function
                    updates = K.function(inputs,[model.total_loss]+model.metrics_tensors,updates=updates,name='train_function_%s'%i,**model._function_kwargs)
                    update_functions.append(updates)
                def F(ins):
                    # 将多个K.function封装为一个单独的函数
                    # 一个K.function就是一次sess.run
                    for f in update_functions:
                        _ = f(ins)
                    return _
                # 最后只需要将model的train_function属性改为对应的函数
                model.train_function = F
    def step1(self, loss, params, old_grads):
        ops = []
        grads = K.gradients(loss, params)
        for p,g,og in zip(params, grads, old_grads[0]):
            ops.append(K.update(og, g))
            ops.append(K.update(p, p - self.lr * g))
        return ops
    def step2(self, loss, params, old_grads):
        ops = []
        grads = K.gradients(loss, params)
        for p,g,og in zip(params, grads, old_grads[0]):
            ops.append(K.update(p, p - 0.5 * self.lr * (g - og)))
        return ops