Keras
Deep Learning
Custom Layers
Neural Networks
Python

How to use keras layers in custom keras layer

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

Using built-in Keras layers inside your own custom layer is a normal and recommended pattern. The important part is to treat those inner layers as child objects of the custom layer so Keras can track their weights, training behavior, and serialization correctly.

Compose Layers Instead Of Rewriting Them

A custom layer does not need to reimplement Dense, Dropout, or LayerNormalization from scratch. In many cases you only want a reusable block that combines existing layers.

A good example is a projection block with normalization and activation.

python
1import tensorflow as tf
2
3
4class ProjectionBlock(tf.keras.layers.Layer):
5    def __init__(self, units: int, rate: float = 0.1, **kwargs):
6        super().__init__(**kwargs)
7        self.dense = tf.keras.layers.Dense(units)
8        self.norm = tf.keras.layers.LayerNormalization()
9        self.dropout = tf.keras.layers.Dropout(rate)
10
11    def call(self, inputs, training=False):
12        x = self.dense(inputs)
13        x = self.norm(x)
14        x = tf.nn.relu(x)
15        return self.dropout(x, training=training)

Here the custom layer is mostly orchestration. Keras still tracks the weights of self.dense and the state of self.norm because they were created as attributes of the layer instance.

Where To Create Child Layers

Create sublayers in __init__ when their configuration is known at construction time. That is the cleanest pattern.

python
1class MLPBlock(tf.keras.layers.Layer):
2    def __init__(self, hidden_units: int, output_units: int, **kwargs):
3        super().__init__(**kwargs)
4        self.hidden = tf.keras.layers.Dense(hidden_units, activation="relu")
5        self.output_layer = tf.keras.layers.Dense(output_units)
6
7    def call(self, inputs):
8        return self.output_layer(self.hidden(inputs))

If a weight shape depends on the input shape and cannot be decided earlier, use build for raw weights you create yourself. Child layers can still often stay in __init__, because Keras builds them lazily when first called.

A Runnable Model Example

python
1inputs = tf.keras.Input(shape=(8,))
2x = ProjectionBlock(16)(inputs)
3outputs = tf.keras.layers.Dense(1)(x)
4model = tf.keras.Model(inputs, outputs)
5model.summary()

The model summary includes the parameters from the child Dense and LayerNormalization layers inside the custom block.

Forward The training Argument When Needed

Any child layer with different training and inference behavior, such as Dropout or BatchNormalization, should receive the training flag from your custom layer.

python
1class ClassifierHead(tf.keras.layers.Layer):
2    def __init__(self, units: int, **kwargs):
3        super().__init__(**kwargs)
4        self.dropout = tf.keras.layers.Dropout(0.2)
5        self.dense = tf.keras.layers.Dense(units)
6
7    def call(self, inputs, training=False):
8        x = self.dropout(inputs, training=training)
9        return self.dense(x)

If you forget to pass training, the child layer may not behave as intended during training or evaluation.

Prefer call For Computation, Not Construction

Do not create new Keras layers inside call on every invocation.

python
1# Bad pattern
2class BadLayer(tf.keras.layers.Layer):
3    def call(self, inputs):
4        dense = tf.keras.layers.Dense(8)
5        return dense(inputs)

That recreates weights each time the layer is called and breaks tracking. The correct pattern is to create the child layer once and reuse it.

Serialization Considerations

If you want to save and reload the custom layer cleanly, implement get_config when the constructor has meaningful arguments.

python
1class ProjectionBlock(tf.keras.layers.Layer):
2    def __init__(self, units: int, rate: float = 0.1, **kwargs):
3        super().__init__(**kwargs)
4        self.units = units
5        self.rate = rate
6        self.dense = tf.keras.layers.Dense(units)
7        self.norm = tf.keras.layers.LayerNormalization()
8        self.dropout = tf.keras.layers.Dropout(rate)
9
10    def call(self, inputs, training=False):
11        x = self.dense(inputs)
12        x = self.norm(x)
13        x = tf.nn.relu(x)
14        return self.dropout(x, training=training)
15
16    def get_config(self):
17        config = super().get_config()
18        config.update({"units": self.units, "rate": self.rate})
19        return config

That makes the layer easier to save as part of a model.

Custom Layer Versus Custom Model

If your object represents a reusable transformation block, subclass Layer. If it represents a whole trainable network with its own top-level behavior, subclass Model. The mechanics are similar, but the intent is different.

Common Pitfalls

The most common mistake is creating child layers inside call, which recreates variables and confuses Keras tracking. Another is forgetting to forward the training flag to layers such as Dropout and BatchNormalization. Developers also sometimes try to manage child-layer weights manually even though Keras already does that when the sublayers are attached as attributes. Finally, custom layers with constructor arguments should implement get_config if model saving matters.

Summary

  • Built-in Keras layers can and should be reused inside custom layers.
  • Create child layers once, usually in __init__, and call them in call.
  • Forward the training flag to layers with training-specific behavior.
  • Avoid constructing new layers inside call.
  • Add get_config when you want robust serialization support.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.