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

作者丨苏剑林

单位丨广州火焰信息科技有限公司

研究方向丨NLP，神经网络

个人主页丨kexue.fm

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

细想之下，不管用什么框架，自定义优化器这个需求可谓真的是小众中的小众。一般而言，对于大多数任务我们都可以无脑地直接上 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] # 优化器带来的权重，在保存模型时会被保存