How to wrap a custom TensorFlow loss function in Keras?

Question

This is my third attempt to get a deep learning project off the ground. I'm working with protein sequences. First I tried TFLearn, then raw TensorFlow, and now I'm trying Keras.

The previous two attempts taught me a lot, and gave me some code and concepts that I can re-use. However there has always been an obstacle, and I've asked questions that the developers can't answer (in the case of TFLearn), or I've simply gotten bogged down (TensorFlow object introspection is tedious).

I have written this TensorFlow loss function, and I know it works:

def l2_angle_distance(pred, tgt):
    with tf.name_scope("L2AngleDistance"):
        # Scaling factor
        count = tgt[...,0,0]
        scale = tf.to_float(tf.count_nonzero(tf.is_finite(count)))
        # Mask NaN in tgt
        tgt = tf.where(tf.is_nan(tgt), pred, tgt)
        # Calculate L1 losses
        losses = tf.losses.cosine_distance(pred, tgt, -1, reduction=tf.losses.Reduction.NONE)
        # Square the losses, then sum, to get L2 scalar loss.
        # Divide the loss result by the scaling factor.
        return tf.reduce_sum(losses * losses) / scale

My target values (tgt) can include NaN, because my protein sequences are passed in a 4D Tensor, despite the fact that the individual sequences differ in length. Before you ask, the data can't be resampled like an image. So I use NaN in the tgt Tensor to indicate "no prediction needed here." Before I calculate the L2 cosine loss, I replace every NaN with the matching values in the prediction (pred) so the loss for every NaN is always zero.

Now, how can I re-use this function in Keras? It appears that the Keras Lambda core layer is not a good choice, because a Lambda only takes a single argument, and a loss function needs two arguments.

Alternately, can I rewrite this function in Keras? I shouldn't ever need to use the Theano or CNTK backend, so it isn't necessary for me to rewrite my function in Keras. I'll use whatever works.

I just looked at the Keras losses.py file to get some clues. I imported keras.backend and had a look around. I also found https://keras.io/backend/. I don't seem to find wrappers for ANY of the TensorFlow function calls I happen to use: to_float(), count_nonzero(), is_finite(), where(), is_nan(), cosine_distance(), or reduce_sum().

Thanks for your suggestions!

Answer 1

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.

Answer 2

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")