how can i implement this custom loss function in tensorflow?

Question

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'),
tf.keras.layers.Dense(100,activation='tanh'),
tf.keras.layers.Dense(50,activation='sigmoid')])

model_s.compile(optimizer='adam',
             loss=loss_fn,
             metrics="accuracy",)


model_s.fit(x_train,y_train,epochs=5,batch_size=512)

I have implemented the function in two ways, but in both cases I encounter an error. case 1:

def loss_fn(y_true,y_pred):
    y=tf.math.argmax(y_true)
    loss=(-(tf.math.log(y_pred[y])+tf.reduce_sum(tf.math.log(1-y_pred[:y]))))
    return tf.reduce_mean(loss)

#example
y_true=tf.constant([0,0,0,0,1])
y_pred=tf.constant([0,0,0,0,.9])
loss_fn(y_true,y_pred)

<tf.Tensor: shape=(), dtype=float32, numpy=0.105360545>

error case 1:

TypeError: Only integers, slices (`:`), ellipsis (`...`), tf.newaxis (`None`) and scalar tf.int32/tf.int64 tensors are valid indices, got <tf.Tensor 'loss_fn/ArgMax:0' shape=(50,) dtype=int64>

case 2:

def loss_fn(y_true,y_pred):
    i=0
    loss=[]
    for pre in y_pred:
        y=tf.math.argmax(y_true[i])
        loss.append(-(tf.math.log(pre[y])+tf.reduce_sum(tf.math.log(1-pre[:y]))))
        i += 1
    return tf.reduce_mean(loss, axis=-1)

#example
y_true=tf.constant([[0,0,0,0,1],[0,0,0,1,0]])
y_pred=tf.constant([[0,0,0,0,.99],[0,0,0,.9,.3]])
loss_fn(y_true,y_pred)
<tf.Tensor: shape=(), dtype=float32, numpy=0.057705436>

error case 2:

ValueError: Tried to convert 'input' to a tensor and failed. Error: The tensor 'Tensor("loss_fn/while/Neg:0", shape=(), dtype=float32)' cannot be accessed here: it is defined in another function or code block. Use return values, explicit Python locals or TensorFlow collections to access it. Defined in: FuncGraph(name=loss_fn_while_body_3814, id=2579287922312); accessed from: FuncGraph(name=train_function, id=2579286968840).

I know the first case may not work for different batch sizes due to its input form, but I have changed the details of the second case a lot, but one thing has always remained the same: showing an error and not working function:(

Answer 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>