Confusion with custom loss for tensorflow keras

Viewed 261

I am trying to follow the many variations of creating a custom loss function for tensorflow.keras

I have successfully created a custom metric which seems to work, and now I would like to use that metric when calculating loss.

Here is the custom metric, which calculates spearman rank correlation between y_true and y_pred...

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import tensorflow_probability as tfp

 class Correlation(keras.metrics.Metric):
    def __init__(self, name="Correlation", **kwargs):
        super(Correlation, self).__init__(name=name, **kwargs)
        self.metric = self.add_weight(name='correlation_01', initializer='zeros')
        
    def update_state(self, y_true, y_pred, sample_weight=None):
        y_true_flat = layers.Flatten()(y_true)
        y_pred_flat = layers.Flatten()(y_pred)
        y_true_rank = tf.cast(tf.argsort(y_true_flat, axis=0, direction="ASCENDING"), 'float32')
        y_pred_rank = tf.cast(tf.argsort(y_pred_flat, axis=0, direction="ASCENDING"), 'float32')
        cov = tfp.stats.covariance(y_true_rank, y_pred_rank, sample_axis=0, event_axis=None)
        std_y_trueR = tfp.stats.stddev(y_true_rank, sample_axis=0, keepdims=False, name="STD_TRUE")
        std_y_predR = tfp.stats.stddev(y_pred_rank, sample_axis=0, keepdims=False, name="STD_PRED")
        #corr = cov/(std_y_trueR * std_y_predR)
        corr = tf.math.divide(cov,tf.math.multiply(std_y_trueR,std_y_predR))
        
        self.metric.assign(corr[0])
        
    def result(self):
        return self.metric
    
    def reset_states(self):
        # reset state of metric at the start of each epoch
        self.metric.assign(0.0)

I don't know how to tie this directly to loss so I thought I would start by following the examples in the docs which suggest creating a function and passing that to compile. So I duplicated the above code and converted the correlation to loss by subtracting it from 1.0

def correlation_loss(y_true, y_pred):
        y_true_flat = layers.Flatten()(y_true)
        y_pred_flat = layers.Flatten()(y_pred)
        y_true_rank = tf.cast(tf.argsort(y_true_flat, axis=0, direction="ASCENDING"), 'float32')
        y_pred_rank = tf.cast(tf.argsort(y_pred_flat, axis=0, direction="ASCENDING"), 'float32')
        cov = tfp.stats.covariance(y_true_rank, y_pred_rank, sample_axis=0, event_axis=None)
        std_y_trueR = tfp.stats.stddev(y_true_rank, sample_axis=0, keepdims=False, name="LOSS_STD_TRUE")
        std_y_predR = tfp.stats.stddev(y_pred_rank, sample_axis=0, keepdims=False, name="LOSS_STD_PRED")
        corr = tf.math.divide(cov,tf.math.multiply(std_y_trueR,std_y_predR))
        loss = tf.math.subtract(1.0,corr[0])
        return loss

I compiled the model with:

model.compile(optimizer='Adam', loss=correlation_loss, metrics=[Correlation()])

However when I attempt to train the model I get a ValueError regarding No gradients provided for any variable

I would love to know what I'm doing wrong. And more specifically, is there a recommended method for using a metric within the loss calculation that does not involve recomputing calculations a second time?

1 Answers

Thank you for pointing me to the other post based on careful reading of that I adjusted my function to insure that each layer was labeled and that seemed to correct the issue.

There is also a problem with my function in that it is not correctly calculating the Spearman Rank correlation. I'm posting the solution to my code here with the rank sections commented out, so it will return standard Pearson Correlation, however this does seem to at least work with the training loop.

def correlation_loss(y_true, y_pred):
    y_true_flat = layers.Flatten(name="Y_TRUE_FLAT")(y_true)
    y_pred_flat = layers.Flatten(name="Y_PRED_FLAT")(y_pred)
    ##-## I can't seem to get the proper ranking for Spearman Correlation
    ##-## Just to have something that functions, I've commented these out for now
    #y_true_rank = tf.cast(tf.argsort(y_true_flat, axis=0, direction="ASCENDING"), 'float32')
    #y_pred_rank = tf.cast(tf.argsort(y_pred_flat, axis=0, direction="ASCENDING"), 'float32')
    cov = tfp.stats.covariance(y_true_flat, y_pred_flat, sample_axis=0, event_axis=None, name="COVARIANCE")
    std_y_trueR = tfp.stats.stddev(y_true_flat, sample_axis=0, keepdims=False, name="LOSS_STD_TRUE")
    std_y_predR = tfp.stats.stddev(y_pred_flat, sample_axis=0, keepdims=False, name="LOSS_STD_PRED")
    corr = tf.math.divide(cov,tf.math.multiply(std_y_trueR,std_y_predR, name="MULT_STDs"), name="CORRELATION")
    loss = tf.math.subtract(1.0,corr[0], name="CORR_LOSS")
    return loss
Related