Tensorflow - range wise regression loss

Question

I am trying to create an efficient loss function for the following problem: The loss is a sum of MAE calculated for each range between the red lines. The blue line is the ground truth, the orange line is a prediction, and the red dots mark the index where the value of the ground truth changes from one to another and close the current value range. Values of inputs are within the [0,1] range. The number of value ranges varies; it can be something between 2-12.

Previously, I made a code with TF map_fn but it was VERY slow:

def rwmae_old(y_true, y_pred):
    y_pred = tf.convert_to_tensor(y_pred)
    y_true = tf.cast(y_true, y_pred.dtype)

    # prepare array
    yt_tmp = tf.concat(
        [tf.ones([len(y_true), 1], dtype=y_pred.dtype) * tf.cast(len(y_true), dtype=y_true.dtype), y_true], axis=-1)
    yt_tmp = tf.concat([yt_tmp, tf.ones([len(y_true), 1]) * tf.cast(len(y_true), dtype=y_true.dtype)], axis=-1)

    # find where there is a change of values between consecutive indices
    ranges = tf.transpose(tf.where(yt_tmp[:, :-1] != yt_tmp[:, 1:]))
    ranges_cols = tf.concat(
        [[0], tf.transpose(tf.where(ranges[1][1:] == 0))[0] + 1, [tf.cast(len(ranges[1]), dtype=y_true.dtype)]], axis=0)
    ranges_rows = tf.range(len(y_true))

    losses = tf.map_fn(
        # loop through every row in the array
        lambda ii:
        tf.reduce_mean(
            tf.map_fn(
                # loop through every range within the example and calculate the loss
                lambda jj:
                tf.reduce_mean(
                    tf.abs(
                        y_true[ii][ranges[1][ranges_cols[ii] + jj]: ranges[1][ranges_cols[ii] + jj + 1]] -
                        y_pred[ii][ranges[1][ranges_cols[ii] + jj]: ranges[1][ranges_cols[ii] + jj + 1]]
                    ),
                ),
                tf.range(ranges_cols[ii + 1] - ranges_cols[ii] - 1),
                fn_output_signature=y_pred.dtype
            )
        ),
        ranges_rows,
        fn_output_signature=y_pred.dtype
    )

    return losses

Today, I created a lazy code that just goes through every example in the batch and checks if values change between indices and, if so, calculates MAE for the current range:

def rwmae(y_true, y_pred):
    (batch_size, length) = y_pred.shape
    losses = tf.zeros(batch_size, dtype=y_pred.dtype)

    for ii in range(batch_size):
        # reset loss for the current row
        loss = tf.constant(0, dtype=y_pred.dtype)
        # set current range start index to 0
        ris = 0

        for jj in range(length - 1):
            if y_true[ii][jj] != y_true[ii][jj + 1]:
                # we found a point of change, calculate the loss in the current range and ...
                loss = tf.add(loss, tf.reduce_mean(tf.abs(y_true[ii][ris: jj + 1] - y_pred[ii][ris: jj + 1])))
                # ... update the new range starting point
                ris = jj + 1
        if ris != length - 1:
            # we need to calculate the loss for the rest of the vector
            loss = tf.add(loss, tf.reduce_mean(tf.abs(y_true[ii][ris: length] - y_pred[ii][ris: length])))
        #replace loss in the proper row
        losses = tf.tensor_scatter_nd_update(losses, [[ii]], [loss])

    return losses

Do you think there is any way to improve its efficiency? Or maybe you think there is a better loss function for the problem?

Answer 1

You can try something like this:

import numpy as np
import tensorflow as tf

def rwmae(y_true, y_pred):
    (batch_size, length) = tf.shape(y_pred)
    losses = tf.zeros(batch_size, dtype=y_pred.dtype)
    for ii in tf.range(batch_size):
        ris = 0
        indices= tf.concat([tf.where(y_true[ii][:-1] != y_true[ii][1:])[:, 0], [length-1]], axis=0)
        ragged_indices = tf.ragged.range(tf.concat([[ris], indices[:-1] + 1], axis=0), indices + 1)
        loss = tf.reduce_sum(tf.reduce_mean(tf.abs(tf.gather(y_true[ii], ragged_indices) - tf.gather(y_pred[ii], ragged_indices)), axis=-1, keepdims=True))
        losses = tf.tensor_scatter_nd_update(losses, [[ii]], [tf.math.divide_no_nan(loss, tf.cast(tf.shape(indices)[0], dtype=tf.float32))])
    return losses

data = np.load('/content/data.npy', allow_pickle=True)
y_pred = data[0:2][0]
y_true = data[0:2][1]

print(rwmae(y_true, y_pred), y_true.shape)