苏剑林作者

Keras之小众需求:自定义优化器

今天我们来看一个小众需求:自定义优化器

细想之下,不管用什么框架,自定义优化器这个需求可谓真的是小众中的小众。一般而言,对于大多数任务我们都可以无脑地直接上 Adam,而调参炼丹高手一般会用 SGD 来调出更好的效果,换言之不管是高手新手,都很少会有自定义优化器的需求。

那这篇文章还有什么价值呢?有些场景下会有一点点作用。比如通过学习 Keras 中的优化器写法,你可以对梯度下降等算法有进一步的认识,你还可以顺带看到 Keras 的源码是多么简洁优雅。

此外,有时候我们可以通过自定义优化器来实现自己的一些功能,比如给一些简单的模型(例如 Word2Vec)重写优化器(直接写死梯度,而不是用自动求导),可以使得算法更快;自定义优化器还可以实现诸如“软 batch”的功能。

Keras优化器

我们首先来看 Keras 中自带优化器的代码,位于:

https://github.com/keras-team/keras/blob/master/keras/optimizers.py

简单起见,我们可以先挑 SGD 来看。当然,Keras 中的 SGD 算法已经把 momentum、nesterov、decay 等整合进去了,这使用起来方便,但不利于学习。所以我稍微简化了一下,给出一个纯粹的 SGD 算法的例子: 

from keras.legacy import interfaces
from keras.optimizers import Optimizer
from keras import backend as K


class SGD(Optimizer):
    """Keras中简单自定义SGD优化器
    """

    def __init__(self, lr=0.01, **kwargs):
        super(SGD, 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')

    @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”,不过我不大清楚是不是就叫这个名字,姑且先这样叫着吧。大概的场景是:假如模型比较庞大,自己的显卡最多也就能跑 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()))

应该也很容易理解吧。如果带有动量的情况,写起来复杂一点,但也是一样的。重点就是引入多一个变量来储存累积梯度,然后引入 cond 来控制是否更新,原来优化器要做的事情,都要在 cond 为 True 的情况下才做(梯度改为累积起来的梯度)。对比原始的 SGD,改动并不大。

“侵入式”优化器

上面实现优化器的方案是标准的,也就是按 Keras 的设计规范来做的,所以做起来很轻松。然而我曾经想要实现的一个优化器,却不能用这种方式来实现,经过阅读源码,得到了一种“侵入式”的写法,这种写法类似“外挂”的形式,可以实现我需要的功能,但不是标准的写法,在此也跟大家分享一下。

原始需求来源于之前的文章从动力学角度看优化算法SGD:一些小启示,里边指出梯度下降优化器可以看成是微分方程组的欧拉解法,进一步可以联想到,微分方程组有很多比欧拉解法更高级的解法呀,能不能用到深度学习中?比如稍微高级一点的有“Heun 方法 [1]”

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

这样就清楚显示出后面这一步实际上是对梯度下降的微调。 

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

经过研究 Keras 模型的训练源码,我发现可以这样写:

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

用法是:

opt = HeunOptimizer(0.1)
opt(model)

model.fit(x_train, y_train, epochs=100, batch_size=32)

其中关键思想在代码中已经注释了,主要是 Keras 的优化器最终都会被包装为一个 train_function,所以我们只需要参照 Keras 的源码设计好 train_function,并在其中插入我们自己的操作。在这个过程中,需要留意到 K.function 所定义的操作相当于一次 sess.run 就行了。

注:类似地还可以实现 RK23、RK45 等算法。遗憾的是,这种优化器缺很容易过拟合,也就是很容易将训练集的 loss 降到很低,但是验证集的 loss 和准确率都很差。

优雅的Keras

本文讲了一个非常非常小众的需求:自定义优化器,介绍了一般情况下 Keras 优化器的写法,以及一种“侵入式”的写法。如果真有这么个特殊需求,可以参考使用。

通过 Keras 中优化器的分析研究,我们进一步可以观察到 Keras 整体代码实在是非常简洁优雅,难以挑剔。

参考文献

[1]. https://en.wikipedia.org/wiki/Heun%27s_method

PaperWeekly
PaperWeekly

推荐、解读、讨论和报道人工智能前沿论文成果的学术平台。

理论Keras优化器
7
相关数据
深度学习技术

深度学习(deep learning)是机器学习的分支,是一种试图使用包含复杂结构或由多重非线性变换构成的多个处理层对数据进行高层抽象的算法。 深度学习是机器学习中一种基于对数据进行表征学习的算法,至今已有数种深度学习框架,如卷积神经网络和深度置信网络和递归神经网络等已被应用在计算机视觉、语音识别、自然语言处理、音频识别与生物信息学等领域并获取了极好的效果。

权重技术

线性模型中特征的系数,或深度网络中的边。训练线性模型的目标是确定每个特征的理想权重。如果权重为 0,则相应的特征对模型来说没有任何贡献。

迭代 技术

模型的权重在训练期间的一次更新。迭代包含计算参数在单个批量数据上的梯度损失。

参数技术

在数学和统计学裡,参数(英语:parameter)是使用通用变量来建立函数和变量之间关系(当这种关系很难用方程来阐述时)的一个数量。

验证集技术

验证数据集是用于调整分类器超参数(即模型结构)的一组数据集,它有时也被称为开发集(dev set)。

梯度下降技术

梯度下降是用于查找函数最小值的一阶迭代优化算法。 要使用梯度下降找到函数的局部最小值,可以采用与当前点的函数梯度(或近似梯度)的负值成比例的步骤。 如果采取的步骤与梯度的正值成比例,则接近该函数的局部最大值,被称为梯度上升。

准确率技术

分类模型的正确预测所占的比例。在多类别分类中,准确率的定义为:正确的预测数/样本总数。 在二元分类中,准确率的定义为:(真正例数+真负例数)/样本总数

批次规模技术

一个批次中的样本数。例如,SGD 的批次规模为 1,而小批次的规模通常介于 10 到 1000 之间。批次规模在训练和推断期间通常是固定的;不过,TensorFlow 允许使用动态批次规模。

过拟合技术

过拟合是指为了得到一致假设而使假设变得过度严格。避免过拟合是分类器设计中的一个核心任务。通常采用增大数据量和测试样本集的方法对分类器性能进行评价。

动量技术

优化器的一种,是模拟物理里动量的概念,其在相关方向可以加速SGD,抑制振荡,从而加快收敛

优化器技术

优化器基类提供了计算梯度loss的方法,并可以将梯度应用于变量。优化器里包含了实现了经典的优化算法,如梯度下降和Adagrad。 优化器是提供了一个可以使用各种优化算法的接口,可以让用户直接调用一些经典的优化算法,如梯度下降法等等。优化器(optimizers)类的基类。这个类定义了在训练模型的时候添加一个操作的API。用户基本上不会直接使用这个类,但是你会用到他的子类比如GradientDescentOptimizer, AdagradOptimizer, MomentumOptimizer(tensorflow下的优化器包)等等这些算法。

暂无评论
暂无评论~