TensorFlow Custom Layers

Custom layers let you define entirely new learnable building blocks for neural networks. When no built-in Keras layer does exactly what you need — a novel attention mechanism, a domain-specific normalization, a research paper's non-standard transformation — you write a custom layer. Custom layers integrate seamlessly with Keras models, get tracked by GradientTape, appear in model.summary(), and save and load with the model.

The Lego Analogy

Keras built-in layers (Dense, Conv2D, LSTM) are standard Lego bricks. They fit together in countless configurations. But some models need a custom-shaped piece that the standard set does not include. You can design that custom brick — your custom layer — and it connects with all the standard pieces perfectly.

Anatomy of a Custom Layer

import tensorflow as tf

class MyLayer(tf.keras.layers.Layer):

    def __init__(self, units, activation=None, **kwargs):
        """
        __init__: Store configuration. Do NOT create weights here.
        **kwargs passes arguments like name, dtype to the parent class.
        """
        super().__init__(**kwargs)
        self.units = units
        self.activation = tf.keras.activations.get(activation)

    def build(self, input_shape):
        """
        build: Create weights. Called automatically the first time
        the layer receives input. input_shape tells you the shape of
        the incoming tensor so you can size your weights correctly.
        """
        self.kernel = self.add_weight(
            name='kernel',
            shape=(input_shape[-1], self.units),
            initializer='glorot_uniform',
            trainable=True
        )
        self.bias = self.add_weight(
            name='bias',
            shape=(self.units,),
            initializer='zeros',
            trainable=True
        )
        super().build(input_shape)   # Mark the layer as built

    def call(self, inputs, training=None):
        """
        call: Define the forward computation. This runs every time
        data passes through the layer. training=None allows
        training-specific behavior (like Dropout).
        """
        output = tf.matmul(inputs, self.kernel) + self.bias
        if self.activation is not None:
            output = self.activation(output)
        return output

    def get_config(self):
        """
        get_config: Returns the configuration so Keras can
        serialize and deserialize this layer (needed for model saving).
        """
        config = super().get_config()
        config.update({'units': self.units,
                       'activation': tf.keras.activations.serialize(self.activation)})
        return config

Using a Custom Layer in a Model

# Custom layers plug into Sequential and Functional API just like built-ins
model = tf.keras.Sequential([
    MyLayer(128, activation='relu', input_shape=(784,), name='custom_dense_1'),
    MyLayer(64, activation='relu', name='custom_dense_2'),
    tf.keras.layers.Dense(10, activation='softmax')  # Mix with built-in layers
])

model.summary()
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])
model.fit(x_train, y_train, epochs=10)

Custom Layer: Learnable Scaling

class ScaledDense(tf.keras.layers.Layer):
    """Dense layer with a learnable per-neuron scale factor."""

    def __init__(self, units, **kwargs):
        super().__init__(**kwargs)
        self.units = units

    def build(self, input_shape):
        self.kernel = self.add_weight(
            'kernel', shape=(input_shape[-1], self.units),
            initializer='glorot_uniform'
        )
        self.bias = self.add_weight(
            'bias', shape=(self.units,), initializer='zeros'
        )
        # Extra learnable scale per output neuron
        self.scale = self.add_weight(
            'scale', shape=(self.units,),
            initializer='ones'   # Start at 1.0 (no scaling initially)
        )
        super().build(input_shape)

    def call(self, inputs):
        output = tf.matmul(inputs, self.kernel) + self.bias
        return output * self.scale   # Element-wise scale

Custom Layer: Squeeze-and-Excitation Block

Squeeze-and-Excitation is a real research technique used in EfficientNet. It weighs each channel of a feature map based on how important that channel is globally — learned during training.

class SqueezeExcitation(tf.keras.layers.Layer):
    """Recalibrates channel-wise feature responses."""

    def __init__(self, reduction_ratio=4, **kwargs):
        super().__init__(**kwargs)
        self.reduction_ratio = reduction_ratio

    def build(self, input_shape):
        channels = input_shape[-1]
        reduced  = max(1, channels // self.reduction_ratio)

        self.squeeze = tf.keras.layers.GlobalAveragePooling2D()
        self.fc1 = tf.keras.layers.Dense(reduced, activation='relu')
        self.fc2 = tf.keras.layers.Dense(channels, activation='sigmoid')
        super().build(input_shape)

    def call(self, inputs):
        # Squeeze: global average pooling → channel vector
        se = self.squeeze(inputs)                # (batch, channels)
        # Excitation: two Dense layers learn channel importance
        se = self.fc1(se)                        # (batch, channels/r)
        se = self.fc2(se)                        # (batch, channels)
        # Reshape for broadcasting
        se = tf.reshape(se, [-1, 1, 1, inputs.shape[-1]])
        # Recalibrate: multiply each channel by its importance score
        return inputs * se

# Use in a CNN
x = tf.keras.layers.Conv2D(64, 3, padding='same', activation='relu')(input_layer)
x = SqueezeExcitation(reduction_ratio=4)(x)

Custom Layer With training Flag

class StochasticDepth(tf.keras.layers.Layer):
    """Randomly skip this layer during training (survival probability)."""

    def __init__(self, survival_prob=0.8, **kwargs):
        super().__init__(**kwargs)
        self.survival_prob = survival_prob

    def call(self, inputs, training=None):
        if not training:
            return inputs   # Always use layer at inference
        # During training: skip layer with probability (1 - survival_prob)
        if tf.random.uniform(()) > self.survival_prob:
            return tf.zeros_like(inputs)   # Skip this layer
        return inputs / self.survival_prob  # Scale to maintain expected value

Non-Trainable State in Custom Layers

class RunningStatsLayer(tf.keras.layers.Layer):
    """Tracks the running mean of activations (non-trainable)."""

    def build(self, input_shape):
        self.running_mean = self.add_weight(
            'running_mean',
            shape=(input_shape[-1],),
            initializer='zeros',
            trainable=False    # Optimizer will NOT update this
        )
        super().build(input_shape)

    def call(self, inputs, training=None):
        if training:
            # Update running mean with exponential moving average
            batch_mean = tf.reduce_mean(inputs, axis=0)
            self.running_mean.assign(0.99 * self.running_mean + 0.01 * batch_mean)
        return inputs   # Pass through unchanged; just tracking stats

Saving Models With Custom Layers

# Save works automatically
model.save('custom_layer_model.keras')

# Loading requires passing the custom class
loaded = tf.keras.models.load_model(
    'custom_layer_model.keras',
    custom_objects={'MyLayer': MyLayer,
                    'SqueezeExcitation': SqueezeExcitation}
)

# Or register the class so Keras finds it automatically
@tf.keras.utils.register_keras_serializable(package='MyPackage')
class MyLayer(tf.keras.layers.Layer):
    ...

# Registered layers load without custom_objects
loaded = tf.keras.models.load_model('custom_layer_model.keras')

Debugging Custom Layers

# Test a custom layer in isolation before using it in a model
layer = SqueezeExcitation(reduction_ratio=4)

test_input = tf.random.normal([2, 8, 8, 64])   # batch=2, H=8, W=8, C=64
output = layer(test_input)

print("Input shape: ", test_input.shape)   # (2, 8, 8, 64)
print("Output shape:", output.shape)       # (2, 8, 8, 64) — same shape
print("Output range: ", output.numpy().min(), "–", output.numpy().max())

Custom layers put the full power of TensorFlow mathematics at your disposal while keeping your code organized and compatible with the Keras ecosystem. Any transformation you can express mathematically can become a reusable, trainable, saveable layer. The next topic covers Autoencoders — a powerful architecture that learns compressed representations of data without any labels, enabling denoising, anomaly detection, and generative modeling.