"Make Keras Cooler!": Layer and Model Reuse Techniques

Layer Reuse

Layer reuse is the simplest. You 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') # Initialize a layer and store it

x = layer(x) # First call
x = layer(x) # Second call
x = layer(x) # Third call

It is important to note that you must initialize the layer first and store it as a variable before calling it to ensure that the repeated calls share the same weights. Conversely, the following code does not share weights:

x = Dense(784, activation='relu')(x) 
x = Dense(784, activation='relu')(x) # Does not share weights with the previous one
x = Dense(784, activation='relu')(x) # Does not share weights with the previous one

Model Reuse

Keras models behave similarly to layers. When called, they can be used in the same way as a layer:

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

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

model = Model(x_in, x) # Build the model


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

x = Dense(784, activation='relu')(x)
x = model(x) # Use the model as if it were a layer

model2 = Model(x_in, x)

Those who have read the Keras source code will understand that the reason a model can be used like a layer is that the Model class itself inherits from the Layer class. Therefore, a model naturally inherits some of the same characteristics as a layer.

Model Cloning

Model cloning is similar to model reuse, except that the new model does not share weights with the original model. In other words, it only preserves the exact same model structure, and updates to the two models are independent. Keras provides a dedicated function for model cloning:

from keras.models import clone_model

model2 = clone_model(model1)

Note that clone_model completely copies the structure of the original model and reconstructs a new one, but it does not copy the weight values. This means that for the same input, the results of model1.predict and model2.predict will be different.

If you want to copy the weights as well, you need to manually use set_weights:

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

Cross-referencing

Cross-referencing here refers to using the weights of an existing layer when defining a new layer. Note that this custom layer might have a completely different function from the old layer; they simply share a specific weight. For example, in BERT, when training the Masked Language Model (MLM), the final fully connected layer that predicts word probabilities shares its weights with the Embedding layer.

A reference implementation is as follows:

class EmbeddingDense(Layer):
    """Operation is identical to Dense, but the kernel uses the 
    embedding matrix from an Embedding layer.
    """
    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,)


# Usage
embedding_layer = Embedding(10000, 128)
x = embedding_layer(x) # Call Embedding layer
x = EmbeddingDense(embedding_layer)(x) # Call EmbeddingDense layer

Extracting Intermediate Layers

Sometimes we need to extract features from intermediate layers of a built model and construct a new model. In Keras, this is also 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 from the Middle

Finally, we come to the most difficult part of this article: splitting a model from the middle. Understanding this also allows for operations like inserting or replacing new layers into an existing model. This requirement might seem unusual, but it has been asked on StackOverflow, indicating it has real value.

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} We might need to replace h_2 with a new input and then connect it to the subsequent layers to build a new model. The function of the new model would be: \text{inputs} \to h_3 \to h_4 \to \text{outputs}

If it is a Sequential model, it is quite simple; you can just iterate through 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)

However, if the model has a complex structure, such as a residual structure that is not a single path, it is not that simple. In fact, this requirement is not inherently difficult because Keras has already written the necessary logic; it just hasn’t provided a public interface. Why do I say this? Because when we call an existing model using code like model(x), Keras essentially rebuilds the entire model from scratch. Since it can rebuild the whole model, rebuilding "half" a model is technically feasible, just lacking an interface. For details, refer to the run_internal_graph function in keras/engine/network.py in the Keras source code.

The logic for fully rebuilding a model is inside the run_internal_graph function, and it is not simple, so we should avoid rewriting it if possible. But if we want to use that logic to disassemble a model from the middle without rewriting it, the only way is to "graft" it: by modifying some attributes of the existing model to trick the run_internal_graph function into thinking that the model’s input layer is an intermediate layer rather than the original input layer. With this idea and a careful reading of the run_internal_graph code, it is not hard to derive the following reference code:

def get_outputs_of(model, start_tensors, input_layers=None):
    """start_tensors is the position where the split begins.
    """
    # Create a new model for this operation
    model = Model(inputs=model.input,
                  outputs=model.output,
                  name='outputs_of_' + model.name)
    # Adaptation for ease of use
    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]
    # Core: Overwrite the model's inputs
    model.inputs = start_tensors
    model._input_layers = [x._keras_history[0] for x in input_layers]
    # Adaptation for ease of use
    if len(input_layers) == 1:
        input_layers = input_layers[0]
    # Organize layers, referenced from Model's run_internal_graph function
    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 # Keep only the layers used
    # Calculate output
    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 logic is actually simple. The truly core code consists of only three lines:

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

That is, overwriting model.inputs and model._input_layers achieves the effect of tricking the model into building from an intermediate layer. The rest is mostly adaptation work, not technical hurdles. The line model._layers = layers keeps only the layers used from the intermediate point onwards, which is mainly for the accuracy of the model’s parameter count statistics; if this part is removed, the model’s parameter count will still appear as large as the original entire model.