Custom Layers and Models

The Sequential and Functional APIs cover most architectures. But when you need residual connections, weight sharing, custom loss functions, or non-standard forward passes, you need to go lower to custom layers and model subclassing. This is where TensorFlow's true flexibility becomes apparent.

Why Create Custom Layers?

•   Reuse novel components: if you invent a new type of layer (a custom attention mechanism, a graph convolution, a domain-specific normalisation), wrapping it in a custom Layer makes it reusable across projects.

•      Encapsulate complex logic: a custom layer can contain its own weights, sub-layers, and forward pass logic, all neatly packaged.

•        Enable SavedModel export: custom layers that implement get_config() can be saved and loaded without access to the original class definition.

Writing Custom Layers with tf.keras

Every custom layer subclasses tf.keras.layers.Layer and implements two methods: build(), which creates the layer's weights, and call(), which defines the forward pass:
import tensorflow as tf

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

"A Dense layer with a learnable scale factor applied to the output."""

def __init__(self, units, **kwargs):

super().__init__(**kwargs)

self.units = units

def build(self, input_shape):

build() is called once when the layer first sees real data

self.W = self.add_weight(

name='kernel', shape=(input_shape[-1], self.units),

initializer='glorot_uniform', trainable=True)

self.b = self.add_weight(

name='bias', shape=(self.units,),

initializer='zeros', trainable=True)

self.scale = self.add_weight(

name='scale', shape=(1,),

initializer='ones', trainable=True)

def call(self, x):

return self.scale * (tf.matmul(x, self.W) + self.b)

def get_config(self):

config = super().get_config()

config.update({'units': self.units})

return config

# Use exactly like any built-in layer

layer = ScaledDense(64)

output = layer(tf.ones((4, 32)))   # (4, 64)

Custom Activation Functions

You can define custom activation functions in two ways: as a simple function or as a full Layer with learnable parameters (like the Parametric ReLU):
import tensorflow as tf

Simple function approach

def swish(x):

"""Swish: x * sigmoid(x) — often outperforms ReLU on deeper networks."""

return x * tf.sigmoid(x)

# Use as activation argument

layer = tf.keras.layers.Dense(64, activation=swish)

# Learnable Parametric ReLU

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

"""PReLU: max(alpha*x, x) where alpha is a learnable parameter."""

def build(self, input_shape):

self.alpha = self.add_weight(

name='alpha', shape=input_shape[1:],

initializer='zeros', trainable=True)

def call(self, x):

return tf.maximum(self.alpha * x, x)

Subclassing tf.keras.Model

Model subclassing gives you complete control over the forward pass. You define layers in __init__() and the computation in call(). This is the most flexible approach and is how PyTorch-style models are written in TensorFlow:

import tensorflow as tf

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

"""A basic residual block: output = F(x) + x."""

def __init__(self, units, **kwargs):

super().__init__(**kwargs)

self.dense1 = tf.keras.layers.Dense(units, activation='relu')

self.dense2 = tf.keras.layers.Dense(units)   # no activation

self.bn = tf.keras.layers.BatchNormalization()

def call(self, x, training=False):

residual = x

x = self.dense1(x)

x = self.dense2(x)

x = self.bn(x, training=training)

return tf.keras.activations.relu(x + residual)  # skip connection

class CustomResNet(tf.keras.Model):

def __init__(self, num_classes, **kwargs):

super().__init__(**kwargs)

self.input_dense = tf.keras.layers.Dense(128, activation='relu')

self.res_block1  = ResidualBlock(128)

self.res_block2  = ResidualBlock(128)

self.classifier  = tf.keras.layers.Dense(num_classes, activation='softmax')

def call(self, x, training=False):

x = self.input_dense(x)

x = self.res_block1(x, training=training)

x = self.res_block2(x, training=training)

return self.classifier(x)

model = CustomResNet(num_classes=10)

model.compile(optimizer='adam',

loss='sparse_categorical_crossentropy',

metrics=['accuracy'])

Building Custom Architectures

Custom architectures are especially useful for multi-input models, multi-output models, and architectures where the same layer is applied to multiple inputs (weight sharing). The Functional API is usually cleaner for these cases, while full subclassing is best for dynamic graphs with conditional logic:

# Functional API example: multi-input model

image_input = tf.keras.Input(shape=(224, 224, 3), name='image')

meta_input  = tf.keras.Input(shape=(10,),       name='metadata')

# Process each input stream

x = tf.keras.layers.Conv2D(32, 3, activation='relu')(image_input)

x = tf.keras.layers.GlobalAveragePooling2D()(x)

m = tf.keras.layers.Dense(16, activation='relu')(meta_input)

# Merge streams

combined = tf.keras.layers.Concatenate()([x, m])

output   = tf.keras.layers.Dense(5, activation='softmax')(combined)

model = tf.keras.Model(inputs=[image_input, meta_input], outputs=output)

model.summary()