"Making Keras Cooler!": Freeform Outputs and Flexible Normalization

“让Keras更酷一些！”：随意的输出和灵活的归一化 苏剑林 · 2019-01-27 · Machine Learning model loss function Keras

Continuing the "Making Keras Cooler!" series, let's make Keras a bit more interesting.

This time we'll dig into Keras loss, metric, weights, and progress bars.

You Don't Have to Have an Output

Normally, we define a Keras model like this:

x_in = Input(shape=(784,))
x = x_in
x = Dense(100, activation='relu')(x)
x = Dense(10, activation='softmax')(x)

model = Model(x_in, x)
model.compile(loss='categorical_crossentropy ',
              optimizer='adam',
              metrics=['accuracy'])
model.fit(x_train, y_train, epochs=5)

This kind of model has a plain input-output structure, and the loss is simply computed from the output. However, for more complex models — such as autoencoders, GANs, or Seq2Seq models — this style is sometimes not flexible enough, since the loss doesn't always have to be a function of the output alone. Fortunately, more recent versions of Keras already support more flexible ways of defining the loss. For example, we can write an autoencoder like this:

x_in = Input(shape=(784,))
x = x_in
x = Dense(100, activation='relu')(x)
x = Dense(784, activation='sigmoid')(x)

model = Model(x_in, x)
loss = K.mean((x - x_in)**2)
model.add_loss(loss)
model.compile(optimizer='adam')
model.fit(x_train, None, epochs=5)

The notable features of the above approach are:

1. When calling compile, no loss is passed in. Instead, the loss is defined separately before compile, and then added to the model via add_loss. This lets us write arbitrarily flexible losses — for instance, a loss that depends on the output of some intermediate layer, on the input, and so on.

1. When calling fit, what used to be the target data is now None, since all inputs and outputs are already being passed in through Input. Readers can also check out my earlier post on Seq2Seq: Playing with Keras Seq2Seq for Automatic Title Generation, where this kind of writing style's convenience is even more apparent.

More Freeform Metrics

Another kind of output is the metric used for monitoring during training. Here, "metric" refers to indicators used to measure model performance, such as accuracy, F1 score, etc. Keras comes with some built-in metrics. As in the example at the beginning, adding these metric names to accuracy into model.compile lets them be displayed dynamically during training.

Of course, you could also define your own custom metrics by mimicking the built-in ones. But the problem is that, in the standard way of defining a metric, a metric is computed as some operation between the "output layer" and the "target value". Yet we often want to monitor the evolution of some special quantity during training — for example, I might want to observe how the output of some intermediate layer changes over time — and in that case, the standard metric definition simply doesn't work.

So what can we do? We can look at the Keras source code and trace through how metrics are implemented, and it turns out that metrics are ultimately defined inside two list. By modifying these two list, we can flexibly display whatever metric we want to observe, for example:

x_in = Input(shape=(784,))
x = x_in
x = Dense(100, activation='relu')(x)
x_h = x
x = Dense(10, activation='softmax')(x)

model = Model(x_in, x)
model.compile(loss='categorical_crossentropy ',
              optimizer='adam',
              metrics=['accuracy'])

# 重点来了
model.metrics_names.append('x_h_norm')
model.metrics_tensors.append(K.mean(K.sum(x_h**2, 1)))

model.fit(x_train, y_train, epochs=5)

The code above shows how to monitor the average norm of an intermediate layer during training. As you can see, this mainly involves two list: model.metrics_names, which is the list of metric names (a list of strings), and model.metrics_tensors, which is the metric tensor itself. As long as you append the quantity you want to display here, it will show up during training. Note, though, that you can only add one scalar at a time.

Flexible Weight Normalization

Sometimes we need to impose some constraint on the weights, commonly a normalization such as L2-norm normalization or spectral normalization, though it could be any other constraint too.

There are generally two ways to implement weight constraints. The first is post-processing: directly and forcibly processing the weights after each gradient descent step, i.e.

\begin{equation}\begin{aligned}&\boldsymbol{\theta} \leftarrow \boldsymbol{\theta} - \varepsilon\nabla_{\boldsymbol{\theta}}L(\boldsymbol{\theta})\\ &\boldsymbol{\theta}\leftarrow constraint(\boldsymbol{\theta})\end{aligned}\end{equation}

Clearly, this kind of processing needs to be written into the optimizer's implementation. In fact, this is exactly what Keras provides out of the box — it's simple to use, you just need to set the kernel_constraint or bias_constraint argument when adding a layer. See details here: https://keras.io/constraints/.

The second approach is pre-processing: we want the weights to be processed before being fed into subsequent layers for computation — that is, we want the constraint to be part of the model rather than part of the optimizer. Keras doesn't natively support this scheme, but we can implement it ourselves.

This is where the elegance of Keras's design really shines through. When building a layer object, Keras splits it into two steps: build and call, where the former is responsible for creating the weights and the latter for performing the computation. By default, these two steps happen together, but we can play a bit of a trick and manually execute them separately.

Below is an implementation of Spectral Normalization based on this idea:

class SpectralNormalization:
    """层的一个包装，用来加上SN。
    """

    def __init__(self, layer):
        self.layer = layer

    def spectral_norm(self, w, r=5):
        w_shape = K.int_shape(w)
        in_dim = np.prod(w_shape[:-1]).astype(int)
        out_dim = w_shape[-1]
        w = K.reshape(w, (in_dim, out_dim))
        u = K.ones((1, in_dim))
        for i in range(r):
            v = K.l2_normalize(K.dot(u, w))
            u = K.l2_normalize(K.dot(v, K.transpose(w)))
        return K.sum(K.dot(K.dot(u, w), K.transpose(v)))

    def spectral_normalization(self, w):
        return w / self.spectral_norm(w)

    def __call__(self, inputs):
        with K.name_scope(self.layer.name):
            if not self.layer.built:
                input_shape = K.int_shape(inputs)
                self.layer.build(input_shape)
                self.layer.built = True
                if self.layer._initial_weights is not None:
                    self.layer.set_weights(self.layer._initial_weights)
        if not hasattr(self.layer, 'spectral_normalization'):
            if hasattr(self.layer, 'kernel'):
                self.layer.kernel = self.spectral_normalization(self.layer.kernel)
            if hasattr(self.layer, 'gamma'):
                self.layer.gamma = self.spectral_normalization(self.layer.gamma)
            self.layer.spectral_normalization = True
        return self.layer(inputs)

Usage:

x = SpectralNormalization(Dense(100, activation='relu'))(x)

In other words, all you need to do is add one line, SpectralNormalization, right after defining the layer. As for the principle behind it, we just need to look at the __call__ part: first, the newly created layer is built=False; then we manually run the build method ourselves, then normalize the original weights and overwrite them with the assignment — that's the line self.layer.kernel = self.spectral_normalization(self.layer.kernel).

Calling Keras's Progress Bar

Lastly, as a fun aside, let's mention Keras's built-in progress bar. In the early days, this progress bar was actually one of the things that drew quite a few new users to Keras. Of course, nowadays a progress bar is nothing new — Python has the excellent tqdm library, which I introduced a long time ago in Two Amazing Python Libraries: tqdm and retry.

That said, if you happen to prefer the look of Keras's progress bar, or don't want to install tqdm separately, you can also call Keras's progress bar in your own code:

import time
from keras.utils import Progbar

pbar = Progbar(100)
for i in range(100):
    pbar.update(i + 1)
    time.sleep(0.1)

It will display progress and the estimated remaining time. If you want to show more information on the progress bar, you can add the value argument when calling update, for example:

import time
from keras.utils import Progbar

pbar = Progbar(100)
for i in range(100):
    pbar.update(i + 1, values=[('something', i - 10)])
    time.sleep(0.1)

One thing to note, though: the value here gets smoothed via a moving average, since this progress bar was mainly designed by Keras for displaying metrics. If you don't want it to update with smoothing, then do:

import time
from keras.utils import Progbar

pbar = Progbar(100, stateful_metrics=['something'])
for i in range(100):
    pbar.update(i + 1, values=[('something', i - 10)])
    time.sleep(0.1)

For more usage details, see here, or check the source code. All in all, its functionality is far less powerful than tqdm's, but as a neat little tool, it's still a decent choice to reach for now and then.

Keras, Endlessly Tinkered With

Once again I've shared some fancy Keras tricks, and I hope they help. Using Keras flexibly and well is a genuinely enjoyable pursuit. Keras may not be the best deep learning framework out there, but it should be the most elegant one (in terms of design/encapsulation) — quite possibly without rival.

Life is short, I use Keras~