Building Custom Layers in Keras

Keras is a tool that helps us build deep learning models easily which gives us many layers like Dense, Conv2Do and LSTM that we can use in our models. But sometimes we want to do something different that these layers can’t do. In such cases we can create our own custom layer where we can write our own code to tell the layer how it should work. Keras makes this easy by letting us create a new class and define what happens inside the layer.

Implementation Custom Layers in Keras

Lets see its implementation step by step:

Step 1: Defining the Custom Layer

This block imports TensorFlow and NumPy then defines a custom layer SimpleDenseLayer. It behaves like a Dense layer creating trainable weights and biases.

• __init__: Initializes the custom layer, accepting a parameter units which determines the number of units in the layer.
• build: Creates the layer’s weights (weights and bias) based on the input shape. These weights are trainable during model training.
• add_weight: A method used to create a weight variable for the layer. It takes shape, initializer and trainable status as arguments.
• call: The forward pass of the layer. It computes the output by multiplying inputs with weights and adding bias.
• get_config: Retrieves the configuration of the layer, allowing it to be saved and reused later. It includes the units parameter used to create the layer.

We can change values and layer according to our needs.

import tensorflow as tf
import numpy as np

class SimpleDenseLayer(tf.keras.layers.Layer):
    def __init__(self, units=32, **kwargs):
        super(SimpleDenseLayer, self).__init__(**kwargs)
        self.units = units

    def build(self, input_shape):
        self.w = self.add_weight(
            shape=(input_shape[-1], self.units),
            initializer='random_normal',
            trainable=True
        )
        self.b = self.add_weight(
            shape=(self.units,),
            initializer='zeros',
            trainable=True
        )

    def call(self, inputs):
        return tf.matmul(inputs, self.w) + self.b

    def get_config(self):
        config = super().get_config()
        config.update({"units": self.units})
        return config

Step 2: Generate Sample Training and Test Data

This block creates random NumPy array for training and test input features and labels. These act as dummy data for model training and prediction.

x_train = np.random.rand(100, 10).astype(np.float32)
y_train = np.random.rand(100, 1).astype(np.float32)
x_test = np.random.rand(10, 10).astype(np.float32)

Step 3: Create and Train the Model Using the Custom Layer

Here, a simple sequential model is built using the custom layer, a ReLU activation and an output layer. The model is compiled with MSE loss and trained for 3 epochs.

model = tf.keras.Sequential([
    SimpleDenseLayer(64),
    tf.keras.layers.ReLU(),
    SimpleDenseLayer(1)
])

model.compile(optimizer='adam', loss='mse')
model.fit(x_train, y_train, epochs=3)

Output:

Step 4: Save the Trained Model

This line saves the entire trained model including architecture, weights and training configuration to an HDF5 file named "my_custom_model.h5".

model.save("my_custom_model.h5")
print("\n Model saved successfully.")

Output:

Model saved successfully.

Step 5: Load the Model with Custom Layer Support

This block reloads the saved model using load_model() and provides the custom layer definition in custom_objects to ensure it can be properly reconstructed.

loaded_model = tf.keras.models.load_model(
    "my_custom_model.h5",
    custom_objects={
        "SimpleDenseLayer": SimpleDenseLayer,
        "mse": tf.keras.losses.MeanSquaredError()
    }
)
print(" Model loaded successfully.")

Output:

Model loaded successfully.

Step 6: Make Predictions with the Loaded Model

Finally this block uses the reloaded model to make predictions on test data and prints the first five outputs to verify successful loading and inference.

predictions = loaded_model.predict(x_test)
print("\n Sample Predictions from Loaded Model:")
print(predictions[:5])

Output:

We can see that our model is working fine and by followig these steps we can build our own Custom Layers in Keras.

You can download the Source code from here: Building Custom Layers in Keras