``Make Keras Cooler!'': Arbitrary Outputs and Flexible Normalization

Output is Optional

Generally, when we define a model in Keras, it looks 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 type of model follows a standard input-output structure where the loss is a function of the output. However, for more complex models like Autoencoders, GANs, or Seq2Seq, this approach is sometimes inconvenient because the loss is not always just a function of the final output. Fortunately, recent versions of Keras support more flexible loss definitions. 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 key characteristics of this approach are:

1. No loss is passed during compile. Instead, the loss is defined before compilation using other methods and added to the model via add_loss. This allows for arbitrary and flexible loss definitions; for instance, the loss can depend on the output of an intermediate layer, the input itself, etc.

2. During fit, the original target data is now None because all necessary inputs and outputs have already been passed through Input layers. Readers can refer to my previous Seq2Seq implementation: “Playing with Keras: Seq2Seq for Automatic Title Generation”. In that example, you can more fully appreciate the convenience of this writing style.

More Arbitrary Metrics

Another type of output is the metric used for observation during training. Metrics here refer to indicators used to measure model performance, such as accuracy or F1 score. Keras has several built-in metrics. Like the accuracy in the first example, adding these metric names to model.compile allows them to be displayed dynamically during training.

Of course, you can also define new metrics by referencing Keras’s built-in ones. However, the problem with the standard metric definition method is that a metric is expected to be a calculation between the “output layer” and the “target value.” Often, we want to observe the changes of special quantities during training—for example, the output changes of a specific intermediate layer. In such cases, standard metric definitions fail.

What can we do? By looking at the Keras source code and tracing its metric-related methods, I found that metrics are actually defined in two lists. By modifying these two lists, we can flexibly display the metrics 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'])

# The key part
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 demonstrates how to observe the change in the average norm of an intermediate layer during training. As you can see, it involves two lists: model.metrics_names is a list of strings representing the names, and model.metrics_tensors is a list of tensors for the metrics. As long as you add the quantities you want to display here, they will appear during training. Note that you can only add one scalar at a time.

Flexible Weight Normalization

Sometimes we need to impose constraints on weights. Common examples include normalization, such as L2 norm normalization, Spectral Normalization, etc.

There are generally two ways to implement weight constraints. The first is post-processing, where the weights are directly manipulated after each gradient descent step: \begin{aligned} &\boldsymbol{\theta} \leftarrow \boldsymbol{\theta} - \varepsilon\nabla_{\boldsymbol{\theta}}L(\boldsymbol{\theta})\\ &\boldsymbol{\theta}\leftarrow \text{constraint}(\boldsymbol{\theta}) \end{aligned} Obviously, this method must be written into the optimizer’s implementation. In fact, Keras’s built-in constraints use this method. Usage is simple: just set the kernel_constraint or bias_constraint parameters when adding a layer. For details, refer to: https://keras.io/constraints/.

The second is pre-processing, where we process the weights before they are substituted into the subsequent layer operations. This means the constraint is part of the model rather than the optimizer. Keras does not natively support this scheme, but we can implement it ourselves.

This is where the brilliance of Keras’s design shines. When creating a layer object, Keras splits it into two steps: build and call. The former is responsible for creating weights, and the latter for the computation. By default, these two parts are executed sequentially, but we can use a “grafting” technique to execute them manually in steps.

Below is an implementation of Spectral Normalization (SN) using this idea:

class SpectralNormalization:
    """A layer wrapper used to add 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)

The usage is:

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

Essentially, you just wrap the defined layer with SpectralNormalization. As for the principle, we only need to observe the __call__ part. First, a newly created layer has built=False. We then manually execute the build method, normalize the original weights, and overwrite them, as seen in the line self.layer.kernel = self.spectral_normalization(self.layer.kernel).

Calling the Keras Progress Bar

Finally, let’s mention a fun little feature: Keras’s built-in progress bar. In the early days, this progress bar attracted many new users to Keras. Of course, progress bars are no longer a novelty; Python has excellent tools like tqdm, which I introduced long ago in “Two Amazing Python Libraries: tqdm and retry”.

However, if you prefer the style of the Keras progress bar or don’t want to install tqdm, you can call Keras’s progress bar in your own scripts:

import time
from keras.utils import Progbar

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

It displays progress and remaining time. If you want to show more content on the progress bar, you can add the values parameter during update:

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)

Note that the values here are subject to moving averages because this progress bar was primarily designed for Keras metrics. If you don’t want it to use moving averages, you can do this:

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 parameters, you can refer to the documentation or the source code. Overall, while it is not as powerful as tqdm, it is a polished tool that is a nice choice for occasional use.

Endless Tinkering with Keras

I have shared some more fancy Keras tricks, and I hope they are helpful to everyone. Using Keras flexibly is quite an enjoyable endeavor. Keras might not be the “best” deep learning framework, but it is likely the most elegant one (wrapper), and quite possibly without equal in that regard.

Life is short, I use Keras.

Original Address: https://kexue.fm/archives/6311

For more details on reprinting, please refer to: Scientific Space FAQ