How to wrap a custom TensorFlow loss function in Keras?

I answered my own question. I'm posting the solution for anyone who may come across this same problem.

I tried using my TF loss function directly in Keras, as was independently suggested by Matias Valdenegro. I did not provoke any errors from Keras by doing so, however, the loss value went immediately to NaN.

Eventually I identified the problem. The calling convention for a Keras loss function is first y_true (which I called tgt), then y_pred (my pred). But the calling convention for a TensorFlow loss function is pred first, then tgt. So if you want to keep a Tensorflow-native version of the loss function around, this fix works:

def keras_l2_angle_distance(tgt, pred):
    return l2_angle_distance(pred, tgt)

<snip>

model.compile(loss = keras_l2_angle_distance, optimizer = "something")

Maybe Theano or CNTK uses the same parameter order as Keras, I don't know. But I'm back in business.


You don't need to use keras.backend, as your loss is directly written in TensorFlow, then you can use it directly in Keras. The backend functions are an abstraction layer so you can code a loss/layer that will work with the multiple available backends in Keras.

You just have to put your loss in the model.compile call:

model.compile(loss = l2_angle_distance, optimizer = "something")