Layers, activations and custom models

Dense, convolution, embedding and recurrent layers, choosing activations deliberately, and writing your own layer and model by subclassing.

Layer	Input shape	Output shape	Typical use
`Dense(n)`	`(batch, features)`	`(batch, n)`	Tabular data, classification heads
`Conv2D(f, k)`	`(batch, h, w, c)`	`(batch, h', w', f)`	Images; `padding="same"` keeps size
`Embedding(v, d)`	`(batch, seq)` ints	`(batch, seq, d)`	Token and category embeddings
`LSTM(units)`	`(batch, steps, f)`	`(batch, units)`	Sequences where order matters
`GlobalAveragePooling2D`	`(batch, h, w, c)`	`(batch, c)`	Replaces a flatten before the head
`LayerNormalization()`	any	same	Transformers and sequence models
`Dropout(p)`	any	same in inference	Regularisation, needs `training=True` at train time

import tensorflow as tf

model = tf.keras.Sequential([
    tf.keras.layers.Input(shape=(28, 28, 1)),
    tf.keras.layers.Conv2D(32, 3, padding="same", activation="relu"),
    tf.keras.layers.MaxPooling2D(),
    tf.keras.layers.Conv2D(64, 3, padding="same", activation="relu"),
    tf.keras.layers.GlobalAveragePooling2D(),
    tf.keras.layers.Dropout(0.3),
    tf.keras.layers.Dense(10),            # logits: no activation
])

model.summary()                            # check the output shape of every stage

Convolution weights are shared across space, so the parameter count depends on kernel size and channel counts, not on image resolution.
Always place a global pooling layer instead of Flatten before the head: it removes the sharp parameter jump and reduces overfitting.
The final layer has no activation when the loss expects logits. Dense(1, activation="sigmoid") plus from_logits=True is a common and silent accuracy loss.

Choosing an activation

# activations as an explicit comparison on the same tiny problem
import numpy as np

def make(activation):
    return tf.keras.Sequential([
        tf.keras.layers.Input(shape=(16,)),
        tf.keras.layers.Dense(32, activation=activation),
        tf.keras.layers.Dense(32, activation=activation),
        tf.keras.layers.Dense(1),
    ])

X = np.random.default_rng(0).normal(size=(2000, 16)).astype("float32")
y = (X[:, 0] * X[:, 1] > 0).astype("float32")

for act in ["relu", "tanh", "gelu", "sigmoid"]:
    m = make(act)
    m.compile(optimizer="adam", loss=tf.keras.losses.BinaryCrossentropy(from_logits=True))
    h = m.fit(X, y, epochs=5, batch_size=64, verbose=0)
    print(act, round(h.history["loss"][-1], 4))

ReLU is the default for feed-forward and convolutional stacks: cheap, no saturation for positive inputs, and it does not vanish on one side.
Sigmoid and tanh saturate at both ends, so stacking them kills gradients. Use them at the output (sigmoid for binary) or inside a gated unit, not as a deep stack's workhorse.
GELU and SiLU (swish) cost a little more compute and often train slightly better; GELU is standard in transformer blocks.
For regression targets with a bounded range, matching the output activation (for example sigmoid scaled to [0, 1]) helps; for unbounded targets, leave the output linear.

Subclassing layers and models

class ScaledDense(tf.keras.layers.Layer):
    """Dense layer with a learnable output scale."""

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

    def build(self, input_shape):
        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.alpha = self.add_weight(
            name="alpha", shape=(self.units,),
            initializer="ones", trainable=True)
        super().build(input_shape)          # must be last

    def call(self, inputs, training=False):
        return tf.nn.relu(inputs @ self.w + self.b) * self.alpha

    def get_config(self):
        return dict(super().get_config(), units=self.units)


class MultiHead(tf.keras.Model):
    def __init__(self, n_classes):
        super().__init__()
        self.backbone = tf.keras.Sequential([
            tf.keras.layers.Dense(64, activation="relu"),
            tf.keras.layers.Dense(32, activation="relu"),
        ])
        self.head = tf.keras.layers.Dense(n_classes)

    def call(self, inputs, training=False):
        return self.head(self.backbone(inputs, training=training))


model = tf.keras.Sequential([
    tf.keras.layers.Input(shape=(16,)),
    ScaledDense(32),
    ScaledDense(16),
    tf.keras.layers.Dense(1),
])
model.summary()

Create weights in build with add_weight, never in __init__: build knows the input shape and runs once.
Call super().build(input_shape) at the end of your own build, otherwise Keras does not mark the layer as built and rebuilds it.
Implement get_config on any custom layer you want to save and reload; without it model.save fails on the custom class.
Subclassed models are harder to serialise and inspect than a functional model. Use tf.keras.Model(inputs, outputs) on tensors unless you genuinely need Python control flow inside call.

⚠️

A custom layer that works in eager mode can still fail inside tf.function. Avoid Python-side effects, use tf.* ops instead of NumPy, and re-test after wrapping the training step in @tf.function.

FAQ

Sequential, functional or subclassing?

Sequential for a plain stack, functional (Model(inputs, outputs)) when you have branches, shared layers or multiple outputs, subclassing only when the forward pass needs Python logic that the functional API cannot express.

How do I know what input shape a layer wants?

Call model.summary() and read the output shape column of the previous layer. For an image model it is (batch, h, w, channels); batch size is always None and comes from the data.

Building a Keras model Debugging, profiling and export formats

Last refreshed 2026-09-18.

Layers, activations and custom models

The layers you will actually use

Choosing an activation

Subclassing layers and models

FAQ

Related