"Making Keras Cooler!": Intermediate Variables, Weight Moving Averages, and Process-Safe Generators

“让Keras更酷一些！”：中间变量、权重滑动和安全生成器 苏剑林 · 2019-04-28 · Machine Learning model optimization Keras

Continuing the "Making Keras Cooler" journey.

Today we'll use Keras to flexibly output arbitrary intermediate variables, perform seamless exponential moving averages of weights, and finally take a look at a process-safe way of writing generators.

First, outputting intermediate variables. When we build custom layers, we sometimes want to inspect intermediate variables. Some of these needs are easy to satisfy — for example, to check the output of some intermediate layer, we just need to save the part of the model up to that layer as a new model. But some needs are much trickier — for example, when using an Attention layer, we might want to inspect the values of the attention matrix, and doing this by building a new model would be very cumbersome. This post gives a simple method that fully satisfies this need.

Next, weight moving averages. Weight moving averaging is an effective way to stabilize and speed up training, and can even improve model performance. Many large models (especially GANs) almost always use weight moving averaging. Generally, weight moving averaging is implemented as part of the optimizer, so it usually requires rewriting the optimizer to add it. This post introduces an implementation of weight moving averaging that can be seamlessly plugged into any Keras model, without needing a custom optimizer.

As for process-safe generators, this comes up because Keras uses multiprocessing when reading from generators, and if the generator itself also involves some multiprocessing operations, this can lead to exceptions. So we need to address this issue. more

Outputting Intermediate Variables

This section takes the basic model

x_in = Input(shape=(784,))
x = x_in

x = Dense(512, activation='relu')(x)
x = Dropout(0.2)(x)
x = Dense(256, activation='relu')(x)
x = Dropout(0.2)(x)
x = Dense(num_classes, activation='softmax')(x)

model = Model(x_in, x)

as an example, and works through, step by step, how to obtain intermediate variables in Keras.

As a New Model

Suppose that after the model has finished training, I want to get the output corresponding to x = Dense(256, activation='relu')(x). I can save the relevant variable when defining the model, and then define a new model:

x_in = Input(shape=(784,))
x = x_in

x = Dense(512, activation='relu')(x)
x = Dropout(0.2)(x)
x = Dense(256, activation='relu')(x)
y = x
x = Dropout(0.2)(x)
x = Dense(num_classes, activation='softmax')(x)

model = Model(x_in, x)
model2 = Model(x_in, y)

Once model has finished training, we can directly use model2.predict to view the corresponding 256-dimensional output. The precondition for this to work is that y must be the output of some layer — it cannot be an arbitrary tensor.

K.function!

Sometimes we define a fairly complex custom layer — a typical example is the Attention layer — and we want to inspect some intermediate variables of the layer, such as the corresponding attention matrix. This becomes tricky, because if we want to use the approach above, we would have to split the original Attention layer into two separate layer definitions. As mentioned earlier, when defining a new Keras model, both the inputs and outputs must be inputs/outputs of Keras layers — they cannot be arbitrary tensors. This means that if we want to inspect several intermediate variables of a layer separately, we would have to keep splitting the layer into more and more sub-layers, which is clearly not very friendly.

Actually, Keras provides an ultimate solution: K.function!

Before introducing K.function, let's write a simple example:

class Normal(Layer):
    def __init__(self, **kwargs):
        super(Normal, self).__init__(**kwargs)
    def build(self, input_shape):
        self.kernel = self.add_weight(name='kernel', 
                                      shape=(1,),
                                      initializer='zeros',
                                      trainable=True)
        self.built = True
    def call(self, x):
        self.x_normalized = K.l2_normalize(x, -1)
        return self.x_normalized * self.kernel

x_in = Input(shape=(784,))
x = x_in

x = Dense(512, activation='relu')(x)
x = Dropout(0.2)(x)
x = Dense(256, activation='relu')(x)
x = Dropout(0.2)(x)
normal = Normal()
x = normal(x)
x = Dense(num_classes, activation='softmax')(x)

model = Model(x_in, x)

In the example above, Normal defines a layer whose output is self.x_normalized * self.kernel, but I want to obtain the value of self.x_normalized after training. This depends on the input, but it is not the output of a layer. So the previous approach won't work here — but with K.function, it's just one line of code:

fn = K.function([x_in], [normal.x_normalized])

The usage of K.function is similar to defining a new model — you need to pass in all the input tensors related to normal.x_normalized — but it does not require the output to be the output of a layer; it can be any arbitrary tensor! The returned fn is an object that behaves like a function, so simply

fn([x_test])

lets us get the x_normalized corresponding to x_test! This is much simpler and more general than defining a new model.

In fact, K.function is one of the basic low-level functions of Keras — it directly wraps the input-output operations of the backend. In other words, when you use TensorFlow as the backend, fn([x_test]) is essentially equivalent to

sess.run(normal.x_normalized, feed_dict={x_in: x_test})

That's why the output of K.function can be any arbitrary tensor — because it's already directly operating on the backend.

Weight Moving Average

Weight moving averaging is an effective way to improve training stability, and it can improve the quality of the solution at almost zero extra cost. Weight moving averaging usually refers to the "Exponential Moving Average," or EMA for short, since the moving average typically uses exponential decay as the weighting scheme. It has been widely adopted by mainstream models, especially GANs. In many GAN papers, we often see descriptions like:

we use an exponential moving average with decay 0.999 over the weight ...

This means the GAN model used EMA. Beyond GANs, ordinary models also use it — for example, QANet: Combining Local Convolution with Global Self-Attention for Reading Comprehension used EMA during training, with a decay rate of 0.9999.

The Form of the Moving Average

The form of the moving average is actually very simple. Suppose each optimizer update is:

\begin{equation}\boldsymbol{\theta}_{n+1} = \boldsymbol{\theta}_n - \Delta \boldsymbol{\theta}_n \end{equation}

Here $\Delta \boldsymbol{\theta}_n$ is the update produced by the optimizer, which can be SGD, Adam, or any other optimizer. The moving average then maintains a new set of variables $\boldsymbol{\Theta}$:

\begin{equation}\boldsymbol{\Theta}_{n+1} = \alpha \boldsymbol{\Theta}_n + (1-\alpha) \boldsymbol{\theta}_{n+1}\end{equation}

where $\alpha$ is a positive constant close to 1, called the "decay rate."

Weight moving averaging is also known as Polyak averaging. Note that, although it looks somewhat similar in form, it is different from momentum acceleration: EMA does not change the trajectory of the original optimizer — whatever path the optimizer originally follows, it still follows the same path — it just maintains an additional set of variables that average over the optimizer's trajectory. Momentum, on the other hand, actually changes the optimizer's trajectory.

To emphasize again: weight moving averaging does not change the direction the optimizer takes; it merely averages points along the optimizer's optimization trajectory, and uses that average as the final model weights.

For more on the theory and effects of weight moving averaging, see A Dynamical Systems View of Optimization Algorithms (Part 4): The Third Stage of GANs.

A Clever Injection-Based Implementation

The key to implementing EMA is figuring out how to introduce a new set of averaged variables on top of the existing optimizer, and how to update these averaged variables after every parameter update. This requires some understanding of the Keras source code and its underlying implementation logic.

Here is a reference implementation:

class ExponentialMovingAverage:
    """对模型权重进行指数滑动平均。
    用法：在model.compile之后、第一次训练之前使用；
    先初始化对象，然后执行inject方法。
    """
    def __init__(self, model, momentum=0.9999):
        self.momentum = momentum
        self.model = model
        self.ema_weights = [K.zeros(K.shape(w)) for w in model.weights]
    def inject(self):
        """添加更新算子到model.metrics_updates。
        """
        self.initialize()
        for w1, w2 in zip(self.ema_weights, self.model.weights):
            op = K.moving_average_update(w1, w2, self.momentum)
            self.model.metrics_updates.append(op)
    def initialize(self):
        """ema_weights初始化跟原模型初始化一致。
        """
        self.old_weights = K.batch_get_value(self.model.weights)
        K.batch_set_value(zip(self.ema_weights, self.old_weights))
    def apply_ema_weights(self):
        """备份原模型权重，然后将平均权重应用到模型上去。
        """
        self.old_weights = K.batch_get_value(self.model.weights)
        ema_weights = K.batch_get_value(self.ema_weights)
        K.batch_set_value(zip(self.model.weights, ema_weights))
    def reset_old_weights(self):
        """恢复模型到旧权重。
        """
        K.batch_set_value(zip(self.model.weights, self.old_weights))

It's very simple to use:

EMAer = ExponentialMovingAverage(model) # 在模型compile之后执行
EMAer.inject() # 在模型compile之后执行

model.fit(x_train, y_train) # 训练模型

After training completes:

EMAer.apply_ema_weights() # 将EMA的权重应用到模型中
model.predict(x_test) # 进行预测、验证、保存等操作

EMAer.reset_old_weights() # 继续训练之前，要恢复模型旧权重。还是那句话，EMA不影响模型的优化轨迹。
model.fit(x_train, y_train) # 继续训练

Looking back at the implementation, the key point is that it introduces a K.moving_average_update operation and inserts it into model.metrics_updates. During training, the model reads and executes all the ops in model.metrics_updates, thereby completing the moving average update.

Process-Safe Generators

Generally speaking, when the training data can't be entirely loaded into memory, or when training data needs to be generated dynamically, we use generator. Typically, the way to write this for Keras models with generator is:

def data_generator():
    while True:
        x_train = something
        y_train = otherthing
        yield x_train, y_train

But if someting or otherthing contains multiprocessing operations of its own, problems can arise. There are two ways to fix this: one is to set the parameter use_multiprocessing=False, worker=0 when calling fit_generator; the other is to write the generator by subclassing the keras.utils.Sequence class.

Official Reference Example

The official documentation for the keras.utils.Sequence class is here. The official docs emphasize:

Sequence are a safer way to do multiprocessing. This structure guarantees that the network will only train once on each sample per epoch which is not the case with generators.

In short, it is safe for multiprocessing, so you can use it with confidence. The official example is as follows:

from skimage.io import imread
from skimage.transform import resize
import numpy as np

# Here, `x_set` is list of path to the images
# and `y_set` are the associated classes.

class CIFAR10Sequence(Sequence):

    def __init__(self, x_set, y_set, batch_size):
        self.x, self.y = x_set, y_set
        self.batch_size = batch_size

    def __len__(self):
        return int(np.ceil(len(self.x) / float(self.batch_size)))

    def __getitem__(self, idx):
        batch_x = self.x[idx * self.batch_size:(idx + 1) * self.batch_size]
        batch_y = self.y[idx * self.batch_size:(idx + 1) * self.batch_size]

        return np.array([
            resize(imread(file_name), (200, 200))
               for file_name in batch_x]), np.array(batch_y)

You just need to define the __len__ and __getitem__ methods according to the required format, and the __getitem__ method directly returns a batch of data.

The bert-as-service Example

I first discovered the necessity of Sequence while experimenting with bert as service. bert as service is a service component built by Han Xiao (肖涵) for quickly obtaining BERT encoding vectors. I once wanted to use it to get character embeddings and then feed them into Keras for training, but I found that training would always freeze up at some point.

After some searching, I confirmed that this was a conflict between the multiprocessing built into Keras's fit_generator and the multiprocessing built into bert-as-service. I'm a bit fuzzy on exactly how they conflict, so I won't dig further into that. But here is a reference solution, which subclasses Sequence to write the generator.

(PS: As far as calling bert as service goes, Han Xiao later provided a coroutine-based version of ConcurrentBertClient, which can replace the original BertClient, so that even with the original generator there would be no problem.)

Keras, a Breath of Fresh Air

In my eyes, Keras is a breath of fresh air among deep learning frameworks, just as Python is a breath of fresh air among programming languages. Doing what you need to do with Keras feels, again and again, like a pleasant experience.