Tensorflow/Keras custom loss function

Question

I want to create a custom loss function in keras.

Let's say I have yTrue and yPred which are tensors (n x m) of true and predicted labels.

Let's call each sample n (that is, each row in yTrue and yPred) yT and yP. Then I want a loss function that computes (yT-yP)^2 when yT[0] == 1, otherwise it will compute (yT[0]-yP[0])^2.

That is: for each sample I always want to calculate the squared error for the first element - but I want to calculate the squared error of the other elements only if the first element of the true label == 1.

How do I do this in a custom loss function?

This is what I have gotten so far:

I need to do this with tensors operations. First I can compute

Y = (yTrue - yPred)^2

Then I can define a masking matrix where the first column is always one, and the others are 1 depending on the value of the first element for each row of yTrue. So I can get something like

1 0 0 0 0 
1 0 0 0 0
1 1 1 1 1
1 1 1 1 1 
1 0 0 0 0 

I can then multiply element wise this matrix with Y and obtain what I want.

However, how do I get in generating the masking matrix? In particular, how do I do the condition "if the first element of the row is 1" in tensorflow/keras?

Maybe there is a better way to do this? Thanks

Answer 1

You can use a conditional switch K.switch in the backend. Something along the lines of:

mse = K.mean(K.square(y_pred - y_true), axis=-1) # standard mse
msep = K.square(y_pred[:,0] - y_true[:,0])
return K.switch(K.equals(y_true[:,0], 1), mse, msep)

Edit for handling per sample condition.