``Making Keras Even Cooler!'': Niche Custom Optimizers

Keras Optimizers

First, let’s look at the built-in optimizer code in Keras, located at:
https://github.com/keras-team/keras/blob/master/keras/optimizers.py

For simplicity, we can start with SGD. Of course, the SGD algorithm in Keras has already integrated momentum, Nesterov, decay, etc. While convenient to use, it is not ideal for learning. Therefore, I have simplified it slightly to provide an 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):
    """Simple custom SGD optimizer in Keras
    """

    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):
        """Main parameter update algorithm
        """
        grads = self.get_gradients(loss, params) # Get gradients
        self.updates = [K.update_add(self.iterations, 1)] # Define set of assignment operators
        self.weights = [self.iterations] # Weights brought by the optimizer, saved with the model
        for p, g in zip(params, grads):
            # Gradient descent
            new_p = p - self.lr * g
            # If there are constraints, apply them to parameters
            if getattr(p, 'constraint', None) is not None:
                new_p = p.constraint(new_p)
            # Add assignment
            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()))

It shouldn’t need much explanation, right? Doesn’t it feel exceptionally simple? Defining an optimizer isn’t such a high-brow task after all.

Implementing “Soft Batch”

Now let’s implement a slightly more complex function, the so-called “soft batch.” I’m not entirely sure if that’s the official name, but let’s call it that for now. The scenario is: suppose the model is quite large, and your GPU can only handle a batch size of 16, but you want to achieve the effect of a batch size of 64. What can you do? One possible solution is to calculate a batch size of 16 each time, cache the gradients, and only update the parameters after 4 batches. That is, gradients are calculated for every small batch, but parameters are updated only once every 4 batches.

Update July 8, 2019: The following implementation is incorrect and does not achieve the expected effect. For a corrected implementation, please see “Keras Gradient Accumulation Optimizer: Trading Time for Results.”

If there is truly a need for this, it can only be solved by modifying the optimizer. Based on the previous SGD, the reference code is as follows:

class7 MySGD(Optimizer):
    """Simple custom SGD optimizer in Keras
    Updates parameters only after a certain number of batches
    """
    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 # How many batches per update

    @interfaces.legacy_get_updates_support
    def get_updates(self, loss, params):
        """Main parameter update algorithm
        """
        shapes = [K.int_shape(p) for p in params]
        sum_grads = [K.zeros(shape) for shape in shapes] # Accumulated gradients for descent
        grads = self.get_gradients(loss, params) # Current batch gradients
        self.updates = [K.update_add(self.iterations, 1)] # Define set of assignment operators
        self.weights = [self.iterations] + sum_grads # Weights saved with the model
        for p, g, sg in zip(params, grads, sum_grads):
            # Gradient descent
            new_p = p - self.lr * sg / float(self.steps_per_update)
            # If there are constraints, apply them to parameters
            if getattr(p, 'constraint', None) is not None:
                new_p = p.constraint(new_p)
            cond = K.equal(self.iterations % self.steps_per_update, 0)
            # Update parameters only when condition is met
            self.updates.append(K.switch(cond, K.update(p, new_p), p))
            # Re-accumulate if condition met, otherwise continue accumulating
            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 to understand. If momentum is involved, the code becomes a bit more complex, but the logic remains the same. The key is to introduce an additional variable to store accumulated gradients and use a cond to control whether to update. Everything the original optimizer does should only happen when cond is True (using the accumulated gradients instead). Compared to the original SGD, the changes are not significant.

“Intrusive” Optimizer

The scheme for implementing optimizers above is standard, following Keras’s design specifications, making it quite easy. However, I once wanted to implement an optimizer that could not be achieved this way. After reading the source code, I derived an “intrusive” approach. This method acts like a “plugin” to implement the required functionality, but it is not a standard implementation. I would like to share it with you here.

The original requirement stems from a previous article “Optimization Algorithms from a Dynamics Perspective (I): From SGD to Momentum Acceleration”, which pointed out that gradient descent optimizers can be viewed as Euler’s method for solving systems of differential equations. This leads to the thought: there are many methods more advanced than Euler’s method for solving differential equations; can they be applied to deep learning? For example, the slightly more advanced “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} where p is the parameter (vector), g is the gradient, and p_i represents the result of p at the i-th iteration. This algorithm requires two steps. Essentially, it performs a standard gradient descent step first (probing), then takes an average based on the probe result to get a more accurate step. Equivalently, it can 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} This clearly shows that the second step is actually a fine-tuning of the gradient descent.

However, implementing such algorithms poses a challenge: gradients must be calculated twice, once for parameter p_i and once for parameter \tilde{p}_{i+1}. In the previous optimizer definition, the get_updates method can only execute one step (corresponding to a single sess.run in the TensorFlow framework; those familiar with TF know it is difficult to implement this requirement with a single sess.run). Therefore, this algorithm cannot be implemented that way. After researching the Keras model training source code, I found it can be written like this:

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

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


class InjectOptimizer(Optimizer):
    """Defines a base class for intrusive optimizers.
    Requires passing in the model and directly modifying the model's training function
    instead of using the optimizer in the conventional way, hence called "intrusive".
    In fact, most of the code below is copied directly from Keras source code:
    https://github.com/keras-team/keras/blob/master/keras/engine/training.py#L497
    Which is the _make_train_function in Keras.
    """
    def get_updates(self, loss, params):
        return []
    def get_grouped_updates(self, loss, params):
        raise NotImplementedError
    def inject(self, model):
        """Pass in the model for injection
        """
        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 Optimizer
    ( 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)
        ]

Usage is as follows:

opt = HeunOptimizer(0.1)
model.compile(loss='mse', optimizer=opt)
opt.inject(model) # This step must be executed

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

The key idea is already noted in the comments. Essentially, Keras optimizers are eventually wrapped into a train_function. Therefore, we only need to design the train_function by referencing the Keras source code and insert our own operations. In this process, it is important to note that an operation defined by K.function is equivalent to a single sess.run.

Note: Similarly, algorithms like RK23 and RK45 can also be implemented. Unfortunately, the performance of such optimizers is not very good.

Elegant Keras

This article discussed a very, very niche requirement: custom optimizers. It introduced the standard way to write Keras optimizers as well as an “intrusive” approach. If you truly have such a special requirement, you can use these as a reference.

Through the analysis and research of optimizers in Keras, we can further observe that the overall Keras code is remarkably concise and elegant, making it hard to find fault with.

Please include the original address when reposting: https://kexue.fm/archives/5879

For more detailed reposting matters, please refer to: “Scientific Space FAQ”