Make a custom loss function in keras

According to the documentation, you can use a custom loss function like this:

Any callable with the signature loss_fn(y_true, y_pred) that returns an array of losses (one of sample in the input batch) can be passed to compile() as a loss. Note that sample weighting is automatically supported for any such loss.

As a simple example:

def my_loss_fn(y_true, y_pred):
    squared_difference = tf.square(y_true - y_pred)
    return tf.reduce_mean(squared_difference, axis=-1)  # Note the `axis=-1`

model.compile(optimizer='adam', loss=my_loss_fn)

Complete example:

import tensorflow as tf
import numpy as np

def my_loss_fn(y_true, y_pred):
    squared_difference = tf.square(y_true - y_pred)
    return tf.reduce_mean(squared_difference, axis=-1)  # Note the `axis=-1`

model = tf.keras.Sequential([
    tf.keras.layers.Dense(8, activation='relu'),
    tf.keras.layers.Dense(16, activation='relu'),
    tf.keras.layers.Dense(1)])

model.compile(optimizer='adam', loss=my_loss_fn)

x = np.random.rand(1000)
y = x**2

history = model.fit(x, y, epochs=10)

In addition, you can extend an existing loss function by inheriting from it. For example masking the BinaryCrossEntropy:

class MaskedBinaryCrossentropy(tf.keras.losses.BinaryCrossentropy):
    def call(self, y_true, y_pred):
        mask = y_true != -1
        y_true = y_true[mask]
        y_pred = y_pred[mask]
        return super().call(y_true, y_pred)

A good starting point is the custom log guide: https://www.tensorflow.org/guide/keras/train_and_evaluate#custom_losses


There are two steps in implementing a parameterized custom loss function in Keras. First, writing a method for the coefficient/metric. Second, writing a wrapper function to format things the way Keras needs them to be.

  1. It's actually quite a bit cleaner to use the Keras backend instead of tensorflow directly for simple custom loss functions like DICE. Here's an example of the coefficient implemented that way:

    import keras.backend as K
    def dice_coef(y_true, y_pred, smooth, thresh):
        y_pred = y_pred > thresh
        y_true_f = K.flatten(y_true)
        y_pred_f = K.flatten(y_pred)
        intersection = K.sum(y_true_f * y_pred_f)
    
        return (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)
    
  2. Now for the tricky part. Keras loss functions must only take (y_true, y_pred) as parameters. So we need a separate function that returns another function.

    def dice_loss(smooth, thresh):
      def dice(y_true, y_pred)
        return -dice_coef(y_true, y_pred, smooth, thresh)
      return dice
    

Finally, you can use it as follows in Keras compile.

# build model 
model = my_model()
# get the loss function
model_dice = dice_loss(smooth=1e-5, thresh=0.5)
# compile model
model.compile(loss=model_dice)