"Making Keras Cooler!": Techniques for Reusing Layers and Models

“让Keras更酷一些！”：层与模型的重用技巧 苏剑林 · 2019-09-29 · Machine Learning model Keras

Today we continue our deep dive into Keras, once again experiencing its unparalleled elegance of design. This time our focus is on "reuse" — mainly the repeated use of layers and models.

Reuse generally aims at two goals: first, sharing weights, meaning that two layers not only perform the same function, but also share weights and are updated synchronously; second, avoiding code duplication — for example, when we've already built a model, and now we want to dismantle it in order to construct some sub-models.

Basics

In fact, Keras has already thought about a lot of this for us, so in many cases, mastering the basic usage is already enough to satisfy most of our needs.

Layer reuse

Reusing a layer is the simplest case: initialize the layer, store it in a variable, and then call it repeatedly:

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

layer = Dense(784, activation='relu') # 初始化一个层，并存起来

x = layer(x) # 第一次调用
x = layer(x) # 再次调用
x = layer(x) # 再次调用

Note that you must first initialize the layer and store it as a variable before calling it, in order to ensure that the repeated calls share the same weights. Conversely, if the code takes the following form, the weights are not shared:

x = Dense(784, activation='relu')(x) 
x = Dense(784, activation='relu')(x) # 跟前面的不共享权重
x = Dense(784, activation='relu')(x) # 跟前面的不共享权重

Model reuse

Keras models behave much like layers — when calling them, you can use the same syntax as with layers, for example:

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

x = Dense(10, activation='softmax')(x)

model = Model(x_in, x) # 建立模型

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

x = Dense(784, activation='relu')(x)
x = model(x) # 将模型当层一样用

model2 = Model(x_in, x)

Readers who have gone through the Keras source code will understand that the reason a model can be used just like a layer is that Model itself is written as a subclass of Layer, so a model naturally inherits some of the same characteristics as a layer.

Model cloning

Model cloning is similar to model reuse, except that the resulting new model does not share weights with the original model. That is, only the exact same model structure is preserved, while the updates to the two models are independent. Keras provides a dedicated function for this, which we can simply call:

from keras.models import clone_model

model2 = clone_model(model1)

Note that clone_model fully replicates the structure of the original model and rebuilds a new model from it, but it does not copy the values of the original model's weights. In other words, given the same input, model1.predict and model2.predict will produce different results.

If we also want to transfer the weights over, we need to manually call set_weights:

model2.set_weights(K.batch_get_value(model1.weights))

Going further

What we discussed above was calling an existing layer or model in an unmodified way, which is relatively simple and already fully supported by Keras. Below, let's look at some more complex examples.

Cross-referencing

Cross-referencing here refers to reusing the weights of an existing layer when defining a new layer — note that this custom layer may have a completely different function from the old layer; the two are purely linked by sharing a particular weight. For example, in Bert, when training the MLM task, the final fully-connected layer that predicts word/character probabilities shares its weights with the Embedding layer.

Here is a reference implementation:

class EmbeddingDense(Layer):
    """运算跟Dense一致，只不过kernel用Embedding层的embedding矩阵
    """
    def __init__(self, embedding_layer, activation='softmax', **kwargs):
        super(EmbeddingDense, self).__init__(**kwargs)
        self.kernel = K.transpose(embedding_layer.embeddings)
        self.activation = activation
        self.units = K.int_shape(self.kernel)[1]

    def build(self, input_shape):
        super(EmbeddingDense, self).build(input_shape)
        self.bias = self.add_weight(name='bias',
                                    shape=(self.units,),
                                    initializer='zeros')

    def call(self, inputs):
        outputs = K.dot(inputs, self.kernel)
        outputs = K.bias_add(outputs, self.bias)
        outputs = Activation(self.activation).call(outputs)
        return outputs
        
    def compute_output_shape(self, input_shape):
        return input_shape[:-1] + (self.units,)

# 用法
embedding_layer = Embedding(10000, 128)
x = embedding_layer(x) # 调用Embedding层
x = EmbeddingDense(embedding_layer)(x) # 调用EmbeddingDense层

Extracting an intermediate layer

Sometimes we need to extract the features of an intermediate layer from an already-built model and construct a new model from them. In Keras, this is likewise a very simple operation:

from keras.applications.resnet50 import ResNet50
model = ResNet50(weights='imagenet')

Model(
    inputs=model.input,
    outputs=[
        model.get_layer('res5a_branch1').output,
        model.get_layer('activation_47').output,
    ]
)

Splitting a model apart from the middle

Finally, we come to the trickiest part of this article: splitting a model apart from the middle. Once you understand this, you'll also be able to insert or replace new layers within an existing model. This requirement may look rather unusual, but as it turns out, someone on Stack Overflow has indeed asked about this], showing that it's a genuinely valuable thing to solve.

Suppose we have an existing model that can be decomposed as

$$\text{inputs}\to h_1 \to h_2 \to h_3 \to h_4 \to \text{outputs}$$

Perhaps we need to replace $h_2$ with a new input, and then connect it to the subsequent layers, in order to construct a new model whose function is:

$$\text{inputs} \to h_3 \to h_4 \to \text{outputs}$$

If it's a Sequential-type model, this is fairly simple — just iterate over all of model.layers to build the new model:

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

for layer in model.layers[2:]:
    x = layer(x)

model2 = Model(x_in, x)

But if the model has a more complex structure — for instance a residual structure that isn't a single straight path from start to end — things aren't so simple. In fact, this requirement isn't inherently difficult; Keras itself has essentially already written the necessary logic, it just doesn't provide a ready-made interface for it. Why do I say this? Because when we call an existing model with code like model(x), Keras is effectively rebuilding that existing model from scratch, from the input all the way to the output. Since it's possible to rebuild the entire model this way, in principle there's no technical obstacle to building "half" a model either — it's just that there's no ready-made interface for it. For details, see the run_internal_graph function of keras/engine/network.py in the Keras source code].

The logic for fully rebuilding a model lives inside the run_internal_graph function, and as you can see, it's not exactly simple — so unless necessary, we're better off not rewriting this code. But if we don't want to rewrite it, and yet still want to call it in order to achieve the effect of splitting a model apart from an intermediate layer, the only option is a bit of "sleight of hand": by modifying certain attributes of the existing model, we trick the run_internal_graph function into believing that the model's input layer is the intermediate layer, rather than the original input layer. With this idea in mind, and after carefully reading through the code of the run_internal_graph function, it's not hard to arrive at the following reference implementation:

def get_outputs_of(model, start_tensors, input_layers=None):
    """start_tensors为开始拆开的位置
    """
    # 为此操作建立新模型
    model = Model(inputs=model.input,
                  outputs=model.output,
                  name='outputs_of_' + model.name)
    # 适配工作，方便使用
    if not isinstance(start_tensors, list):
        start_tensors = [start_tensors]
    if input_layers is None:
        input_layers = [
            Input(shape=K.int_shape(x)[1:], dtype=K.dtype(x))
            for x in start_tensors
        ]
    elif not isinstance(input_layers, list):
        input_layers = [input_layers]
    # 核心：覆盖模型的输入
    model.inputs = start_tensors
    model._input_layers = [x._keras_history[0] for x in input_layers]
    # 适配工作，方便使用
    if len(input_layers) == 1:
        input_layers = input_layers[0]
    # 整理层，参考自 Model 的 run_internal_graph 函数
    layers, tensor_map = [], set()
    for x in model.inputs:
        tensor_map.add(str(id(x)))
    depth_keys = list(model._nodes_by_depth.keys())
    depth_keys.sort(reverse=True)
    for depth in depth_keys:
        nodes = model._nodes_by_depth[depth]
        for node in nodes:
            n = 0
            for x in node.input_tensors:
                if str(id(x)) in tensor_map:
                    n += 1
            if n == len(node.input_tensors):
                if node.outbound_layer not in layers:
                    layers.append(node.outbound_layer)
                for x in node.output_tensors:
                    tensor_map.add(str(id(x)))
    model._layers = layers # 只保留用到的层
    # 计算输出
    outputs = model(input_layers)
    return input_layers, outputs

Usage:

from keras.applications.resnet50 import ResNet50
model = ResNet50(weights='imagenet')

x, y = get_outputs_of(
    model,
    model.get_layer('add_15').output
)

model2 = Model(x, y)

The code is a bit long, but the underlying logic is actually quite simple — the truly essential part is just three lines:

model.inputs = start_tensors
model._input_layers = [x._keras_history[0] for x in input_layers]
outputs = model(input_layers)

That is, by overriding the model's model.inputs and model._input_layers, we can trick the model into building itself starting from the intermediate layer. Most of the rest is adaptation work rather than anything technically deep, and the line model._layers = layers simply retains only the layers actually used starting from the intermediate layer — this is purely to ensure the accuracy of the reported parameter count; if you remove this part, the model's parameter count would still show as that of the entire original model.

Summary

Keras is, without question, the most delightful deep learning framework to work with — at least so far, in terms of the readability of its model code, there is none better. Readers might bring up PyTorch — and indeed, PyTorch has plenty of merits — but in terms of readability, I don't think it measures up to Keras.

In the process of digging deep into Keras, I've not only marveled at the deep and elegant programming skill of its authors, but I even feel that my own programming ability has improved quite a bit as a result. Indeed, many of my Python programming techniques were learned from reading the Keras source code.