Custom Layers — When Built-Ins Don't Cover It

~13 min · custom-layer, build, add-weight

Level 0Level 0

0 XP0/78 lessons0/17 achievements

0/100 XP to next level100 XP to go0% complete

Three methods, no fewer

When built-in layers don't cover your architecture, subclass keras.layers.Layer. Three methods to implement:

__init__ — store config (units, activation, etc.). Do NOT create weights here.
build(input_shape) — create weights here. Called once on first forward pass when input shape is finally known. Use self.add_weight(...) — never raw tf.Variable in Keras 3.
call(inputs) — forward computation.

For serialization (saving/loading), also implement get_config() — return a dict that from_config() can use to reconstruct the layer.

Keras 3 breaking change: use self.add_weight(), not tf.Variable. Raw tf.Variable objects are no longer auto-tracked by Keras 3's variable tracking. Using them inside a layer makes weights invisible to the optimizer — your model trains nothing and you'll wonder why.

Code

CustomDense — 정석 패턴·python

import tensorflow as tf
from keras import layers

class CustomDense(layers.Layer):
    def __init__(self, units=32, activation=None, **kwargs):
        super().__init__(**kwargs)
        self.units = units
        self.activation = tf.keras.activations.get(activation)

    def build(self, input_shape):
        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)

    def call(self, inputs):
        output = tf.matmul(inputs, self.kernel) + self.bias
        if self.activation is not None:
            output = self.activation(output)
        return output

    def get_config(self):
        config = super().get_config()
        config.update({
            "units": self.units,
            "activation": tf.keras.activations.serialize(self.activation),
        })
        return config

# Use in a Sequential model
model = tf.keras.Sequential([
    tf.keras.Input(shape=(784,)),
    CustomDense(256, activation='relu'),
    CustomDense(10),
])

External links

Making new layers and models guide

Exercise

Implement a ScaleLayer custom layer that has a single trainable scalar weight scale initialized to 1.0, and returns inputs * scale. Verify that after training on data where the target is y = 3 * x, the learned scale converges near 3.0.

Progress

Progress is local-only — sign in to sync across devices.

← PreviousEssential Layers and Activation Functions Next →Quiz · 5 questions

Spotted a bug or have feedback on this page?Report an Issue

Comments 0

🔔 Reply notifications (sign in)

No comments yet — be the first.