Learn AI Create the Future — with Pranay
Main site

TensorFlow and Keras

Custom layers, losses and metrics

When the built-in parts stop fitting your problem, Keras lets you write your own layer, your own definition of wrong, and your own scorecard.

Read these first

On this page 5
  1. Why they exist
  2. How it works
  3. A real example you have seen
  4. Remember this
  5. What to learn next

One lesson, three depths. Pick the one that fits you today — you can switch any time.

Beginner — No maths. Plain English.

Keras ships ready-made layers, losses and metrics — and clean extension points for writing your own when the ready-made ones do not fit.

Think of furnishing a house. Most rooms are fine with shop-bought furniture: standard beds, standard cupboards. But one room has a slanted ceiling, and no shop sells a cupboard for it. You call a carpenter and describe the exact shape you need.

Custom layers, losses and metrics are the carpenter option. Same house, same rooms — one part built to your measurements.

Why they exist

Three different things go custom, for three different reasons:

  • A custom layer — when the transformation you need between inputs and outputs is not in the catalogue. Maybe your data needs a learned scaling no standard layer provides.
  • A custom loss — when the built-in definition of "wrong" does not match the real cost of being wrong. A shop that under-stocks loses a sale; over-stocking loses only shelf space. Those mistakes are not equal, but standard losses treat them equally.
  • A custom metric — when the score humans care about is not accuracy. "How often were we within 2 rupees of the real price?" is a business question no built-in metric answers.

How it works

custom layer:    inputs  →  [ your maths + learnable knobs ]  →  outputs
custom loss:     (truth, prediction)  →  one number: "how bad"
custom metric:   (truth, prediction)  →  running score across batches

Each is a small recipe you hand to Keras. The layer joins the model. The loss and metric plug into compile, exactly where the built-in names went.

A real example you have seen

Food delivery apps predict arrival times. Telling you 30 minutes and arriving in 45 makes you angry; telling you 45 and arriving in 30 makes you pleased. The two errors are the same size and completely different in cost. Teams building these systems write a loss that punishes late-side mistakes harder — a custom definition of wrong.

Remember this

  • Layer = custom transformation with learnable values inside the model.
  • Loss = custom definition of "how wrong", used for learning.
  • Metric = custom scorecard for humans, not used for learning.

What to learn next

Developer — Code and libraries.

Setup

bash
pip install tensorflow numpy

Outputs verified with TensorFlow 2.21 (Keras 3.12), CPU, seed 0. Learned weight values will drift slightly across versions.

A custom layer and a custom loss, trained

custom_parts.py
import numpy as np
import keras
from keras import layers, ops

keras.utils.set_random_seed(0)


class ScaledDense(layers.Layer):
    """A Dense layer that also learns one overall output scale."""

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

    def build(self, input_shape):
        # build() runs once, when the layer first sees its input shape
        self.w = self.add_weight(shape=(input_shape[-1], self.units),
                                 initializer="glorot_uniform", trainable=True)
        self.b = self.add_weight(shape=(self.units,),
                                 initializer="zeros", trainable=True)
        self.scale = self.add_weight(shape=(), initializer="ones",
                                     trainable=True)

    def call(self, inputs):
        return self.scale * (ops.matmul(inputs, self.w) + self.b)


def underprediction_loss(y_true, y_pred):
    """Squared error, but predicting too LOW costs 3x more than too high."""
    diff = y_true - y_pred
    weight = ops.where(diff > 0, 3.0, 1.0)
    return ops.mean(weight * ops.square(diff))


# Demand forecasting: yesterday's sales -> today's sales
x = np.array([[10.], [20.], [30.], [40.], [50.], [60.]], dtype="float32")
y = np.array([[12.], [22.], [33.], [41.], [55.], [63.]], dtype="float32")

model = keras.Sequential([keras.Input(shape=(1,)), ScaledDense(1)])
model.compile(optimizer=keras.optimizers.Adam(0.5), loss=underprediction_loss)
model.fit(x, y, epochs=300, verbose=0)

preds = model.predict(x, verbose=0)
print("weights learned:", [round(float(v), 3) for v in
                           [model.layers[0].w[0, 0], model.layers[0].b[0],
                            model.layers[0].scale]])
print("true 33 predicted as:", round(float(preds[2, 0]), 1))
Output
weights learned: [0.909, 1.201, 1.155]
true 33 predicted as: 32.9

The walkthrough

__init__ vs build vs call — the three-part contract of every layer. __init__ stores configuration (how many units). build creates weights, and runs only once, the first time data arrives — that is how the layer knows input_shape without you declaring it. call is the forward computation, run on every batch.

self.add_weight(...) registers a trainable value. Registration is the point: anything added this way is found by the optimizer, saved with the model, and listed in summary(). A plain attribute holding a tensor would train nothing.

ops.matmul, ops.where, ops.mean come from keras.ops — the Keras 3 operations namespace that works on every backend. Writing tf. inside layers still works on the TensorFlow backend, but locks the layer to it.

The loss is a plain function with the fixed signature (y_true, y_pred) → scalar-per-sample. The asymmetry lives in one where: under-predictions (truth above prediction) get weight 3. Passing the function object — not a string — to compile is all the registration needed.

A custom metric

Metrics accumulate across batches, so they are classes with state, not plain functions:

within_tolerance.py
import numpy as np
import keras
from keras import ops


class WithinTolerance(keras.metrics.Metric):
    """Fraction of predictions landing within +/- tolerance of the truth."""

    def __init__(self, tolerance=2.0, name="within_tolerance", **kwargs):
        super().__init__(name=name, **kwargs)
        self.tolerance = tolerance
        self.hits = self.add_weight(shape=(), initializer="zeros", name="hits")
        self.total = self.add_weight(shape=(), initializer="zeros", name="total")

    def update_state(self, y_true, y_pred, sample_weight=None):
        close = ops.abs(y_true - y_pred) <= self.tolerance
        self.hits.assign_add(ops.sum(ops.cast(close, "float32")))
        self.total.assign_add(ops.cast(ops.size(y_true), "float32"))

    def result(self):
        return self.hits / self.total


metric = WithinTolerance(tolerance=2.0)
metric.update_state(np.array([12., 22., 33.]), np.array([11., 30., 32.5]))
print("within 2 units:", round(float(metric.result()), 3))
Output
within 2 units: 0.667

update_state runs per batch and adds to counters; result turns the counters into the final score. Passing metrics=[WithinTolerance()] to compile makes it appear in fit logs like any built-in.

Common mistakes

Creating weights in call. New weights every batch, nothing trains, and Keras raises an error about variables created inside a traced function. Weights belong in build (or __init__ if shapes are known).

A loss with flipped arguments. The signature is (y_true, y_pred) in that order. Swap them in an asymmetric loss and you punish the opposite mistake. Symmetric losses hide the bug entirely — test your loss on hand-made numbers where the two orders differ.

Using a metric as a loss. Metrics like the one above use comparisons, which have no useful gradient. Training on them leaves the model frozen. Losses must be smooth; metrics may be blunt.

Forgetting get_config before saving. A model containing custom pieces reloads only if each piece can describe its constructor arguments. Add get_config() returning {**super().get_config(), "units": self.units} — or reloading raises. Details in saving and loading Keras models.

Try it yourself

Flip the loss: make over-prediction cost 3x instead. Retrain and compare the prediction for the 40. input against the run above — the model should now aim low. One changed line, visible behaviour change.

What to learn next

Researcher — Mathematics and papers.

Asymmetric losses, formally

The custom loss above is a special case of the linex/weighted-quadratic family. With residual $r = y - \hat{y}$:

$$ \ell(r) = w(r)\, r^2, \qquad w(r) = \begin{cases} c & r > 0 \ 1 & r \le 0 \end{cases} $$

Symbols: $y$ — the truth; $\hat{y}$ — the prediction; $r$ — the residual, positive when the model under-predicts; $c$ — the asymmetry factor (3 in the code).

The minimiser of expected loss shifts away from the conditional mean: for squared loss the optimal prediction is $\mathbb{E}[y|x]$, but under asymmetric weighting it moves toward the costly side. In the limit of the pinball (quantile) loss $\ell_\tau(r) = r(\tau - \mathbf{1}[r<0])$, the optimal prediction is exactly the $\tau$-quantile of $y|x$ — the principled version of "aim high" (Koenker and Bassett 1978, Regression quantiles). If your real goal is a quantile, use pinball loss rather than tuning $c$ by feel.

The gradient requirement: backpropagation needs $\partial \ell / \partial \hat{y}$ to exist and be non-zero almost everywhere. The where above is piecewise-smooth — fine. Hard 0/1 comparisons (accuracy, the tolerance metric) have zero gradient almost everywhere, which is the formal reason metrics cannot serve as losses; the usable trick is a smooth surrogate, the same argument that replaced 0-1 classification loss with crossentropy (Bartlett, Jordan, McAuliffe 2006, Convexity, classification, and risk bounds).

Layer lifecycle cost

Deferred building (build on first call) exists so layers compose without shape declarations — shape inference propagates through the model graph once, at first trace, costing O(number of layers) and nothing per step thereafter. The registration machinery (add_weight) is what ties into the pytree/variable tracking that optimizers traverse; unregistered tensors are invisible to model.trainable_variables, which is the single mechanism behind the "nothing trains" failure mode.

References

  • Koenker and Bassett (1978), Regression quantiles.
  • Varian (1975), A Bayesian approach to real estate assessment — origin of the linex asymmetric loss.
  • Bartlett, Jordan, McAuliffe (2006), Convexity, classification, and risk bounds.

What to learn next