``Making Keras Cooler!'': Intermediate Variables, Weight Averaging, and Safe Generators

As a New Model

Suppose that after the model is trained, I want to obtain the output corresponding to x = Dense(256, activation=’relu’)(x). I can save the corresponding variable while defining the model and then redefine 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)

After completing the training of model, you can directly use model2.predict to view the corresponding 256-dimensional output. The prerequisite for this is that y must be the output of a certain layer; it cannot be just any arbitrary tensor.

K.function!

Sometimes we define a relatively complex layer, a typical example being an Attention layer. We hope to view some intermediate variables of the layer, such as the corresponding Attention matrix. This becomes quite troublesome. If we want to use the previous method, we would have to split the original Attention layer into two separate layers because, as mentioned before, when defining a new Keras model, the inputs and outputs must be the inputs and outputs of Keras layers, not arbitrary tensors. Consequently, if you want to view multiple intermediate variables of a layer, you would have to continuously split the layer into multiple layers, which is clearly not user-friendly.

In fact, 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 above example, Normal defines a layer where the output is self.x_normalized * self.kernel. However, I want to obtain the value of self.x_normalized after training is complete. This value depends on the input and is not the output of a layer. Thus, the previous method cannot be used, but with K.function, it is 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 must pass in all input tensors related to normal.x_normalized, but it does not require the output to be a layer output—it allows any arbitrary tensor! The returned fn is an object with function-like capabilities, so you only need:

fn([x_test])

to obtain 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 fundamental functions of the Keras backend. It directly encapsulates the backend’s input and output operations. In other words, when using TensorFlow as the backend, fn([x_test]) is equivalent to:

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

Therefore, the output of K.function is allowed to be any tensor because it is essentially operating directly on the backend.

Format of Moving Average

The format of the moving average is actually very simple. Assume each update of the optimizer is: \boldsymbol{\theta}_{n+1} = \boldsymbol{\theta}_n - \Delta \boldsymbol{\theta}_n Here, \Delta \boldsymbol{\theta}_n is the update brought by the optimizer, which can be any type like SGD, Adam, etc. The moving average maintains a new set of variables \boldsymbol{\Theta}: \boldsymbol{\Theta}_{n+1} = \alpha \boldsymbol{\Theta}_n + (1-\alpha) \boldsymbol{\theta}_{n+1} where \alpha is a positive constant close to 1, called the “decay rate.”

Weight moving averaging is also called Polyak averaging. Note that although it is somewhat similar in form, it is different from momentum acceleration: EMA does not change the trajectory of the original optimizer. That is, however the original optimizer moved, it still moves the same way; it just maintains a new set of variables to average the trajectory of the original optimizer. Momentum acceleration, on the other hand, changes the trajectory of the original optimizer.

To emphasize again: Weight moving averaging does not change the direction of the optimizer; it simply takes the points on the optimizer’s optimization trajectory, averages them, and uses that as the final model weights.

For more on the principles and effects of weight moving averaging, you can further refer to the article “Optimization Algorithms from a Dynamical Perspective (IV): The Third Stage of GANs”.

A Clever Injection Implementation

The key to implementing EMA is how to introduce a new set of average variables on top of the original optimizer and execute the update of these average variables after each parameter update. This requires a certain understanding of Keras source code and its implementation logic.

A reference implementation is provided below:

class ExponentialMovingAverage:
    """Apply exponential moving average to model weights.
    Usage: Use after model.compile and before the first training;
    Initialize the object, then execute the inject method.
    """
    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):
        """Add update operators to 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):
        """Initialize ema_weights to be consistent with the original model.
        """
        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):
        """Backup original weights, then apply average weights to the model.
        """
        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):
        """Restore the model to the old weights.
        """
        K.batch_set_value(zip(self.model.weights, self.old_weights))

The usage is very simple:

EMAer = ExponentialMovingAverage(model) # Execute after model.compile
EMAer.inject() # Execute after model.compile

model.fit(x_train, y_train) # Train the model

After training is complete:

EMAer.apply_ema_weights() # Apply EMA weights to the model
model.predict(x_test) # Perform prediction, validation, saving, etc.

EMAer.reset_old_weights() # Before continuing training, restore old weights.
# Again, EMA does not affect the model's optimization trajectory.
model.fit(x_train, y_train) # Continue training

Looking back at the implementation process, the main point is the introduction of the K.moving_average_update operation and its insertion into model.metrics_updates. During the training process, the model reads and executes all operators in model.metrics_updates, thereby completing the moving average.

Official Reference Example

The official introduction to the keras.utils.Sequence class is here. The official documentation emphasizes:

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 multi-processing and can be used with confidence. The example provided by the official documentation 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 format. The __getitem__ method directly returns one batch of data.

bert-as-service Example

I first discovered the necessity of Sequence while experimenting with bert-as-service. bert-as-service is a service component developed by Han Xiao for quickly obtaining BERT encoding vectors. I once wanted to use it to get character vectors and pass them into Keras for training, but I found that the training would always get stuck after a while.

After searching, it was confirmed to be a conflict between the multi-processing of Keras’s fit_generator and the built-in multi-processing of bert-as-service. I won’t go into the details of the conflict, but a reference solution was provided here, which uses the Sequence class to write the generator.

(PS: Regarding calling bert-as-service, Han Xiao later provided a coroutine version ConcurrentBertClient, which can replace the original BertClient, so that even with a primitive generator, there will be no issues.)