Keras
custom cost function
machine learning
deep learning
neural networks

Can I use my own cost function in keras?

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

Yes, Keras lets you define your own loss function, and doing so is often the right move when built-in losses do not match your business metric or modeling constraint. The important part is that the custom loss must be expressed with TensorFlow operations so gradients can still flow through the graph during training.

Write a Simple Function-Based Loss

The smallest custom loss is a Python function that takes y_true and y_pred and returns a tensor of loss values.

python
1import tensorflow as tf
2from tensorflow import keras
3
4
5def weighted_mse(y_true, y_pred):
6    error = tf.square(y_true - y_pred)
7    weights = tf.where(y_true > 0.5, 2.0, 1.0)
8    return tf.reduce_mean(error * weights)
9
10
11model = keras.Sequential([
12    keras.layers.Dense(16, activation="relu"),
13    keras.layers.Dense(1)
14])
15
16model.compile(optimizer="adam", loss=weighted_mse)

Keras will call this function during training just like it calls built-in losses such as mean squared error or binary crossentropy.

Use TensorFlow Ops, Not Plain Python Math

The loss must operate on tensors, not on plain Python scalars or NumPy arrays created inside the function. That is how TensorFlow tracks gradients.

Good pattern:

python
def custom_loss(y_true, y_pred):
    return tf.reduce_mean(tf.abs(y_true - y_pred))

Bad pattern:

  • converting tensors to NumPy inside the loss
  • using Python loops when vectorized tensor ops are available
  • returning a plain Python number instead of a tensor

If the loss breaks differentiability or leaves the TensorFlow graph, training may fail or silently behave poorly.

Parameterize a Custom Loss

If the loss needs configuration, wrap it in a factory function.

python
1import tensorflow as tf
2
3
4def asymmetric_loss(over_penalty=2.0):
5    def loss(y_true, y_pred):
6        error = y_pred - y_true
7        penalty = tf.where(error > 0, over_penalty, 1.0)
8        return tf.reduce_mean(tf.square(error) * penalty)
9    return loss

Then compile with:

python
model.compile(optimizer="adam", loss=asymmetric_loss(over_penalty=3.0))

This is a clean way to reuse the same loss idea with different business penalties.

Subclass keras.losses.Loss for Reusability

For more structure, especially in larger codebases, define a custom loss class.

python
1import tensorflow as tf
2from tensorflow import keras
3
4
5class WeightedMAE(keras.losses.Loss):
6    def __init__(self, positive_weight=2.0, name="weighted_mae"):
7        super().__init__(name=name)
8        self.positive_weight = positive_weight
9
10    def call(self, y_true, y_pred):
11        weights = tf.where(y_true > 0.5, self.positive_weight, 1.0)
12        return tf.reduce_mean(tf.abs(y_true - y_pred) * weights)

This approach is useful when you want clean serialization logic or a loss object that carries its own configuration.

Test the Loss Before Training

Custom losses can compile successfully and still be wrong mathematically. Before launching a training run, check:

  • output shape and dtype
  • loss value on easy hand-checked examples
  • whether gradients are finite
  • whether the loss decreases on a tiny overfit test

A short sanity check helps catch sign errors, wrong reductions, or accidentally reversed arguments.

Think About Serialization

If you save and reload the model, Keras needs to know how to reconstruct the custom loss. A common pattern is:

python
1model.save("model.keras")
2
3loaded = keras.models.load_model(
4    "model.keras",
5    custom_objects={"weighted_mse": weighted_mse, "WeightedMAE": WeightedMAE}
6)

If you forget custom_objects, model loading can fail even though training worked correctly earlier.

Common Pitfalls

  • Writing the loss with NumPy or plain Python operations that break gradient tracking.
  • Returning a Python scalar instead of a TensorFlow tensor.
  • Skipping sanity tests and discovering a mathematical bug only after long training runs.
  • Forgetting to register custom losses when reloading a saved model.
  • Making the loss overly complex when a weighted or transformed built-in loss would already solve the problem.

Summary

  • Keras supports fully custom loss functions.
  • Use TensorFlow tensor operations so gradients can be computed correctly.
  • Start with a simple function and move to a subclass only if the project needs more structure.
  • Test the loss on small examples before expensive training.
  • Remember to provide custom loss objects again when loading saved models.

Course illustration
Course illustration

All Rights Reserved.