I know that is better avoid loop in Keras custom loss function, but I think I have to do it. The problem is the following: I'm trying to implement a loss function that compute a loss value for multiple bunches of data and then aggregate this values in an unique value.

For example I have 6 data entry, so in my Keras loss I'll have 6 y_true and 6 y_pred. I want to compute 2 loss value: one for the first 3 elements and one for the last 3 elements.

Example of hypothetical code:

def custom_loss(y_true, y_pred): 
    start_range = 0
    losses = []
    for index in range(0,2):
        end_range = start_range + 3
        y_true_bunch = y_true[start_range:end_range]
        y_pred_bunch = y_pred[start_range:end_range]
        loss_value = ...some processing on bunches...
        start_range = end_range
    final_loss = ...aggregate loss_value...
    return final_loss

Is it possible to achieve something like this? I need to process the whole dataset and calculate loss for multiple bunches and then aggregate all bunches value in a single value.


1 Answer 1


Assuming 2 partitions since, for index in range(0,2) and a fixed bunch size of 3 since, start_range + 3

The shape of your y_true should be (nb_samples, 2), where the first vector of 2nd dimension is the ground_truth and the other is partition_index

import numpy as np
import tensorflow as tf
import tensorflow.keras.backend as K
from tensorflow.keras.losses import mean_squared_error

y_true = tf.Variable(np.array([[1.5, 0],[1.2, 0],[1.3, 0],[1.6, 1],[3.0, 1],[2.25, 1]]), dtype=tf.float32)
y_pred = tf.Variable(np.array([[1.35],[1.24],[1.69],[1.55],[1.24],[1.69]]), dtype=tf.float32)

def get_two_splits(y_true):
    child_y = y_true[:, 1]
    child_y = tf.expand_dims(child_y, 1)
    y_true = y_true[:, 0]
    y_true = tf.expand_dims(y_true, 1)
    return y_true, child_y

def some_processing_on_bunches(y_all):
    y_true, y_pred = get_two_splits(y_all)
    return mean_squared_error(y_true, y_pred)

def custom_loss(y_true, y_pred):
    n = K.shape(y_true)[0]
    y_true, partitions = get_two_splits(y_true)
    partitions = tf.reshape(partitions ,[n, ])
    partitions = tf.cast(partitions, tf.int32)

    y_elements = tf.dynamic_partition(tf.concat([y_true, y_pred], 1), partitions, 2)
    losses = tf.vectorized_map(some_processing_on_bunches, tf.stack(y_elements))
    return tf.reduce_mean(losses, axis=1)

K.eval(custom_loss(y_true, y_pred))
>>>array([0.05873336, 1.1379    ], dtype=float32)

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.