How Can I Implement This Custom Loss Function In Tensorflow?

My loss function: I want to implement the above loss function for the following model: model_s=tf.keras.Sequential([ tf.keras.layers.Dense(100, input_shape=(50,),activation='tanh'

Solution 1:

Instead of trying to index into your predictions, just use your true labels (which are in one-hot format) as a mask to get the i-th prediction (you can accomplish this with a multiply and sum across rows). Then use your argmax to create a sequence mask for the second part.

Try this:

import tensorflow as tf


def custom_loss_fn(y_true, y_pred, num_labels=5):
    idx = tf.math.argmax(y_true, 1)
    msk = tf.cast(tf.sequence_mask(idx, num_labels), tf.float32)
    # 1st part
    fst = -tf.math.log(tf.math.reduce_sum(y_true * y_pred, 1))
    # 2nd part
    snd = tf.math.reduce_sum(tf.math.log(1.0 - y_pred * msk, 1))
    return tf.math.reduce_mean(fst + snd)

Test 1:

y_true = tf.constant([[0, 0, 0, 0, 1]], tf.float32)
y_pred = tf.constant([[0, 0, 0, 0, 0.9]])

custom_loss_fn(y_true, y_pred)
# <tf.Tensor: shape=(), dtype=float32, numpy=0.105360545>

Test 2:

y_true = tf.constant([[0, 0, 0, 0, 1], [0, 0, 0, 1, 0]], tf.float32)
y_pred = tf.constant([[0, 0, 0, 0, 0.99], [0, 0, 0, 0.9, 0.3]])

custom_loss_fn(y_true, y_pred)
# <tf.Tensor: shape=(), dtype=float32, numpy=0.057705436>