Creating Custom Layers in Keras: A Step-by-Step Guide

Learn how to create and custom layers in your neural network

Introduction

If you’ve been playing around with TensorFlow/Keras for a while, you probably know that it comes with a big bag of pre-built layers: Dense, Conv2D, LSTM, Dropout—you name it. For many projects, these are all you’ll ever need.

But what if you have a weird, experimental idea for a neural network layer? Or maybe you just want to wrap up some logic you keep reusing in different models. That’s where custom layers come into play.

In this post, I’ll walk you through how to build your own Keras layer from scratch. Don’t worry—it’s not as scary as it sounds. We’ll go step by step, with examples along the way.

What Exactly Is a Custom Layer?

A custom layer is just like any other Keras layer… except you make it yourself. Think of it as baking your own bread instead of buying a loaf from the store. Keras gives you the framework, and you fill in the recipe.

Under the hood, custom layers are created by subclassing tf.keras.layers.Layer. That means you define how the layer should behave (the math it does on inputs) and, if needed, what trainable weights it should have.

Why bother? Three main reasons:

• Flexibility: Add operations that aren’t included in standard layers.

• Reusability: Package your own logic and reuse it across models.

• Experimentation: Try out new neural network ideas quickly.

Visualizing the Lifecycle of a Custom Layer

Here’s the rough “flow” of how a Keras custom layer works:

Inputs → __init__ → build → call → Outputs

• __init__: Define static stuff (non-trainable params).

• build: Create trainable weights (runs once per input shape).

• call: Apply the computation every time data passes through.

Think of it like a restaurant:

• __init__ = writing the menu

• build = stocking the kitchen with ingredients

• call = cooking the dish every time someone orders

Step 1: Imports

First, make sure you’ve got TensorFlow installed:

pip install tensorflow

Then import the usual import classes:

import tensorflow as tf
from tensorflow import keras
from keras import layers

---

Step 2: Subclass Layer

The starting point is always subclassing layers.Layer:

class MyCustomLayer(layers.Layer):
    def __init__(self, **kwargs):
        super(MyCustomLayer, self).__init__(**kwargs)

At this point, our layer doesn’t do much—but we’ve got the skeleton in place.

---

Step 3: Add Trainable Weights (Optional)

If your layer needs its own weights (like a scaling factor, bias, etc.), you define them inside build(). This method runs the first time your layer sees data, so it knows the input shape.

Example: let’s add a single trainable number called scale_factor.

def build(self, input_shape):
    self.scale_factor = self.add_weight(
        name=’scale_factor’,
        shape=(),
        initializer=tf.keras.initializers.Constant(1.0),
        trainable=True
    )

---

Step 4: Define the Computation (call)

This is the heart of your layer. call() tells Keras what to do when data flows through.

def call(self, inputs):
    return inputs * self.scale_factor

Here, every input gets multiplied by that trainable scale factor.

---

Step 5: Make It Serializable (Optional)

If you want your layer to be savable and reloadable, implement get_config():

def get_config(self):
    config = super(MyCustomLayer, self).get_config()
    return config

---

Putting It All Together

Here’s the full layer definition:

class MyCustomLayer(layers.Layer):
    def __init__(self, **kwargs):
        super(MyCustomLayer, self).__init__(**kwargs)

    def build(self, input_shape):
        self.scale_factor = self.add_weight(
            name=’scale_factor’,
            shape=(),
            initializer=tf.keras.initializers.Constant(1.0),
            trainable=True
        )

    def call(self, inputs):
        return inputs * self.scale_factor

    def get_config(self):
        return super(MyCustomLayer, self).get_config()

---

Step 6: Use It in a Model

Now let’s drop it into a simple model:

model = keras.Sequential([
    layers.Input(shape=(10,)),
    MyCustomLayer(),
    layers.Dense(5, activation=’relu’),
    layers.Dense(1)
])

Compile it:

model.compile(optimizer=’adam’, loss=’mse’, metrics=[’mae’])

And train on some toy data:

import numpy as np

x_train = np.linspace(0, 10, 1000).reshape(-1, 1)
y_train = np.sin(x_train)

model.fit(x_train, y_train, epochs=10, batch_size=32, verbose=1)

---

Step 7: Save and Load

When saving, you’ll need to tell Keras about your custom layer:

model.save(’custom_model.h5’)

loaded_model = keras.models.load_model(
    ‘custom_model.h5’,
    custom_objects={’MyCustomLayer’: MyCustomLayer}
)

---

Another Example: Squaring Layer

Here’s a simpler custom layer that just squares every input element. No trainable weights needed:

class SquareLayer(layers.Layer):
    def __init__(self, **kwargs):
        super(SquareLayer, self).__init__(**kwargs)

    def build(self, input_shape):
        pass  # no weights to define

    def call(self, inputs):
        return tf.square(inputs)

    def get_config(self):
        return super(SquareLayer, self).get_config()

Use it in a model:

model = keras.Sequential([
    layers.Input(shape=(10,)),
    SquareLayer(),
    layers.Dense(1)
])

---

Tips for Working With Custom Layers

• Put weights in build() (not __init__) so they adapt to input shapes.

• Always test with some dummy inputs before plugging into a big model.

• Use TensorFlow ops (tf.*) instead of raw NumPy—this keeps your layer compatible with GPUs/TPUs.

• Watch your shapes. Debugging shape mismatches is 90% of building layers.

---

Wrapping Up

And that’s it—you’ve just built your own custom Keras layers!

To recap, the recipe is:

1. Subclass Layer

2. Define weights in build() (if needed)

3. Write the computation in call()

4. (Optional) Add get_config() for serialization

Once you’ve got this down, you can create pretty much any building block you want—attention layers, weird activations, or just quick hacks for experiments.

If you want to go deeper, check out the official Keras docs on custom layers.

Happy model-building! 😀