"Making Keras Cooler!": A Niche Topic — Custom Optimizers

“让Keras更酷一些！”：小众的自定义优化器 苏剑林 · 2018-09-08 · Machine Learning Keras optimizer

Continuing from the earlier post "Making Keras Cooler!": Delicate Layers and Fancy Callbacks.

Today we'll look at a rather niche need: custom optimizers.

Thinking about it, no matter which framework you use, writing a custom optimizer is truly a niche need among niche needs. For most tasks, we can just throw Adam at the problem without thinking too hard, while the real tuning wizards tend to get better results out of SGD. Either way—novice or expert—hardly anyone actually needs to write their own optimizer.

So what's the point of this post then? Well, it turns out to be a bit useful in a few scenarios. For instance, by studying how optimizers are implemented in Keras, you can deepen your understanding of gradient descent and related algorithms, and along the way get a glimpse of just how clean and elegant the Keras source code is. Beyond that, sometimes we might want to customize an optimizer to implement our own special functionality—for example, rewriting the optimizer for a simple model (like Word2Vec) by hard-coding the gradient instead of relying on autodiff, which can speed things up; or a custom optimizer can be used to implement something like "soft batching."

Keras Optimizers

Let's first look at the code for the built-in optimizers in Keras, located at:

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

For simplicity, let's start with SGD. Of course, Keras's built-in SGD implementation already bundles in momentum, Nesterov, decay, and so on, which is convenient to use but not so great for learning purposes. So I've simplified it a bit, giving a bare-bones example of a pure SGD algorithm:

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()))

No need to explain that, right? Doesn't it feel surprisingly simple? Defining an optimizer isn't such a lofty, mysterious thing after all~

Implementing "Soft Batching"

Now let's implement something a little more involved: so-called "soft batching"—though I'm not entirely sure that's the official name, so let's just call it that for now. The rough scenario is this: suppose your model is fairly large, and your GPU can only handle a batch size of 16 at most, but you want the effect of a batch size of 64. What can you do? One possible approach is to compute with batch size 16 each time, cache the gradients, and only update the parameters after 4 batches. In other words, every small batch computes a gradient, but the parameters are only updated once every 4 batches.

Update, July 8, 2019: The implementation below is actually flawed and doesn't achieve the intended effect. For a corrected implementation, see Keras Gradient Accumulation Optimizer: Trading Time for Effect.

If you genuinely need this, the only way to achieve it is by modifying the optimizer. Building on the SGD example above, here's some reference code:

class7 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()))

This should also be easy enough to follow. If momentum is involved, the code gets a bit more complex, but the idea is the same. The key point is to introduce an extra variable to store the accumulated gradient, and then use a cond to control whether an update happens—everything the original optimizer used to do now only happens when cond is True (and the gradient used is the accumulated one). Compared to the original SGD, the changes are actually quite small.

An "Invasive" Optimizer

The approach to implementing optimizers above is the standard one, i.e., it follows Keras's own design conventions, which is why it's so painless. However, there was once an optimizer I wanted to implement that couldn't be done this way. After digging through the source code, I came up with what I'd call an "invasive" approach—more like a "hack" or "plug-in"—which let me achieve what I needed, even though it's not the standard way of doing things. Let me share it here.

The original motivation came from an earlier post, Viewing Optimization Algorithms Through the Lens of Dynamics (I): From SGD to Momentum Acceleration, which pointed out that gradient-descent-based optimizers can be seen as Euler's method applied to a system of differential equations. This naturally leads to the thought: there are plenty of methods for solving differential equations that are more sophisticated than Euler's method—can any of those be applied to deep learning? For example, a somewhat more advanced method is "Heun's method":

$$\begin{aligned}\tilde{p}_{i+1} =& p_i + \epsilon g(p_i)\\ p_{i+1} =& p_i + \frac{1}{2}\epsilon \big[g(p_i)+g(\tilde{p}_{i+1})\big] \end{aligned}$$

Here $p$ is the parameter (vector), $g$ is the gradient, and $p_i$ denotes the result of the $i$-th iteration of $p$. This algorithm requires two steps: roughly speaking, an ordinary gradient descent step is taken first (as a "scouting" step), and then the result of that scouting step is averaged with the original gradient to produce a more precise step. This can equivalently be rewritten as:

$$\begin{aligned}\tilde{p}_{i+1} =& p_i + \epsilon g(p_i)\\ p_{i+1} =& \tilde{p}_{i+1} + \frac{1}{2}\epsilon \big[g(\tilde{p}_{i+1}) - g(p_i)\big] \end{aligned}$$

which makes it clear that the second step is really just a fine-tuning correction on top of gradient descent.

The difficulty in implementing this kind of algorithm, though, is that it requires computing the gradient twice: once with respect to parameter $p_i$, and once with respect to parameter $\tilde{p}_{i+1}$. But the get_updates method in the optimizer definitions we saw earlier can only execute a single step (in TF terms, this corresponds to a single sess.run—and anyone familiar with TF knows it's very hard to satisfy this requirement within just one sess.run). So this kind of algorithm can't be implemented that way. After digging into the source code for Keras's model training, I found that it could be done like this:

#! -*- coding: utf-8 -*-

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

class InjectOptimizer(Optimizer):
    """定义注入式优化器的基类
        需要传入模型，直接修改模型的训练函数，而不按常规流程使用优化器，所以称为“侵入式”
        其实下面的大部分代码，都是直接抄自keras的源码：
        https://github.com/keras-team/keras/blob/master/keras/engine/training.py#L497
        也就是keras中的_make_train_function函数。
    """
    def get_updates(self, loss, params):
        return []
    def get_grouped_updates(self, loss, params):
        raise NotImplementedError
    def inject(self, model):
        """传入模型做注入
        """
        if not hasattr(model, 'train_function'):
            raise RuntimeError('You must compile your model before using it.')
        model._check_trainable_weights_consistency()
        if model.train_function is None:
            inputs = (model._feed_inputs +
                      model._feed_targets +
                      model._feed_sample_weights)
            if model._uses_dynamic_learning_phase():
                inputs += [K.learning_phase()]
            with K.name_scope('training'):
                train_functions = []
                with K.name_scope(model.optimizer.__class__.__name__):
                    grouped_training_updates = self.get_grouped_updates(
                        params=model._collected_trainable_weights,
                        loss=model.total_loss)
                for i, updates in enumerate(grouped_training_updates[1:]):
                    f = K.function(
                        inputs,
                        [model.total_loss],
                        updates=updates,
                        name='train_function_%s' % (i + 1),
                        **model._function_kwargs)
                    train_functions.append(f)
                # Gets loss and metrics. Updates weights at each call.
                first_updates = (model.updates +
                                 grouped_training_updates[0],
                                 model.metrics_updates)
                first_train = K.function(
                    inputs,
                    [model.total_loss] + model.metrics_tensors,
                    updates=first_updates,
                    name='train_function',
                    **model._function_kwargs)
                def F(inputs):
                    R = first_train(inputs)
                    for f in train_functions:
                        f(inputs)
                    return R
                model.train_function = F

class HeunOptimizer(InjectOptimizer):
    """Heun优化器
    ( https://en.wikipedia.org/wiki/Heun%27s_method )
    """
    def __init__(self, lr, **kwargs):
        super(HeunOptimizer, self).__init__(**kwargs)
        with K.name_scope(self.__class__.__name__):
            self.lr = K.variable(lr, name='lr')
    def get_updates_1(self, loss, params, cache_grads):
        updates = []
        grads = self.get_gradients(loss, params)
        for p, g, cg in zip(params, grads, cache_grads):
            updates.append(K.update(cg, g))
            updates.append(K.update(p, p - self.lr * g))
        return updates
    def get_updates_2(self, loss, params, cache_grads):
        updates = []
        grads = self.get_gradients(loss, params)
        for p, g, cg in zip(params, grads, cache_grads):
            updates.append(K.update(p, p - 0.5 * self.lr * (g - cg)))
        return updates
    def get_grouped_updates(self, loss, params):
        cache_grads = [K.zeros(K.int_shape(p)) for p in params]
        return [
            self.get_updates_1(loss, params, cache_grads),
            self.get_updates_2(loss, params, cache_grads)
        ]

Used like so:

opt = HeunOptimizer(0.1)
model.compile(loss='mse', optimizer=opt)
opt.inject(model) # 必须执行这步

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

The key idea is already noted in the comments in the code above: ultimately, all Keras optimizers get wrapped into a train_function, so all we need to do is design our own train_function following Keras's own source design, and insert our own operations into it. The important thing to keep in mind during this process is that whatever operation is defined via K.function is equivalent to a single sess.run.

Note: Similarly, one could implement algorithms like RK23, RK45, and so on. Unfortunately, this kind of optimizer doesn't actually perform very well in practice~

The Elegance of Keras

This post covered a very, very niche need: custom optimizers. We looked at how optimizers are normally written in Keras, as well as an "invasive" style of implementation. If you ever genuinely run into such a special requirement, feel free to use this as a reference.

Through this study of how optimizers are implemented in Keras, we can once again appreciate just how clean and elegant the overall Keras codebase is—it's hard to find fault with it~