Keras - Customized Layer

Keras - Custom Layers

While tf.keras.layers provides a comprehensive set of built-in layers, you might encounter situations where you need a layer with specific, non-standard functionality. This could be for implementing a novel research idea, creating a layer with a unique stateful behavior, or simply combining operations in a reusable way.

tf.keras allows you to create your own custom layers by subclassing the base tf.keras.layers.Layer class. Once defined, your custom layer can be used seamlessly within Sequential or Functional models just like any built-in layer.

Structure of a Custom Layer

To create a custom layer, you typically need to implement three key methods within your subclass:

1. __init__(self, **kwargs): (Constructor) Use this to perform initialization independent of input shape. Define any non-weight attributes or hyperparameters specific to your layer here. Always call super().__init__(**kwargs).
2. build(self, input_shape): This method is called automatically the first time the layer is used, once the shape of the input tensor(s) is known. This is the standard place to create the layer’s trainable weights using the self.add_weight() method. The input_shape argument provides the shape of the input(s). Ensure you call super().build(input_shape) at the end.
3. call(self, inputs, **kwargs): This method defines the layer’s forward pass logic – the core computation performed by the layer. It takes the input tensor(s) as arguments and returns the output tensor(s). This is where you implement the layer’s transformation using TensorFlow operations or calls to other layers.
4. get_config(self) (Optional but Recommended): To allow your layer (and models containing it) to be saved and loaded correctly, implement get_config. This method should return a dictionary containing the constructor arguments needed to recreate the layer instance. Also include the config from the parent class.

Example: A Simple Fully Connected Layer

Let’s create a basic custom dense (fully connected) layer from scratch to illustrate the process. This layer will perform the operation: output = activation(dot(input, kernel) + bias).

import tensorflow as tf
from tensorflow.keras import layers
from tensorflow.keras import activations # Import activations module

class SimpleDense(layers.Layer):
    """A basic custom Dense layer."""

    def __init__(self, units, activation=None, **kwargs):
        super().__init__(**kwargs)
        self.units = units
        # Get the activation function from the identifier (string or function)
        self.activation = activations.get(activation)

    def build(self, input_shape):
        # input_shape is a TensorShape object (e.g., (batch_size, input_dim))
        input_dim = input_shape[-1]

        # Add trainable weight variable for the kernel
        self.kernel = self.add_weight(
            name='kernel',
            shape=(input_dim, self.units),
            initializer='glorot_uniform', # Common initializer for dense layers
            trainable=True
        )

        # Add trainable weight variable for the bias
        self.bias = self.add_weight(
            name='bias',
            shape=(self.units,),
            initializer='zeros',
            trainable=True
        )

        # Ensure the parent class's build method is called
        super().build(input_shape)

    def call(self, inputs):
        # Define the forward pass logic
        # Use tf.matmul for matrix multiplication
        z = tf.matmul(inputs, self.kernel) + self.bias

        # Apply activation function if specified
        if self.activation is not None:
            return self.activation(z)
        return z

    def compute_output_shape(self, input_shape):
        # Optional method to compute output shape symbolically
        # Keras can often infer this, but defining it can be helpful
        return (input_shape[0], self.units)

    def get_config(self):
        # Make the layer serializable
        config = super().get_config()
        config.update({
            'units': self.units,
            'activation': activations.serialize(self.activation) # Serialize activation
        })
        return config

Using the Custom Layer

Now you can use SimpleDense just like a built-in Keras layer:

# Example Usage
model = tf.keras.Sequential([
    tf.keras.Input(shape=(16,)), # Input layer
    SimpleDense(units=32, activation='relu'), # Use the custom layer
    layers.Dropout(0.1), # Can mix with standard layers
    SimpleDense(units=8, activation='softmax') # Use the custom layer again
])

# Compile and train as usual
# model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
model.summary()

Running model.summary() will show your custom layer (simple_dense) listed alongside standard layers.

# Example Summary Output:
# Model: "sequential"
# _________________________________________________________________
#  Layer (type)                Output Shape              Param #
# =================================================================
#  simple_dense (SimpleDense)  (None, 32)                544
#
#  dropout (Dropout)           (None, 32)                0
#
#  simple_dense_1 (SimpleDen   (None, 8)                 264
#  se)
#
# =================================================================
# Total params: 808 (3.16 KB)
# Trainable params: 808 (3.16 KB)
# Non-trainable params: 0 (0.00 Byte)
# _________________________________________________________________

When to Create Custom Layers

• Implementing novel layer types from research papers.
• Creating layers with internal state that needs managing differently from standard weights.
• Encapsulating a complex, reusable block of operations.
• When existing layers don’t offer the exact required functionality or flexibility.

Subclassing tf.keras.layers.Layer provides a powerful mechanism for extending Keras and tailoring it to specific research or application needs.