Check differentiability of custom loss function before training

[Frightera]

Feature request was made in keras-team/tf-keras#52.

This is a very common mistake when users define custom loss function / classes which are not differentiable which leads getting None gradients in fitting process. This check makes it easier for users to interpret the problem.

Example usage:

import tensorflow as tf

X = tf.constant([-7.0, -4.0, -1.0])
y = tf.constant([3.0, 6.0, 9.0])

model = tf.keras.Sequential([
  tf.keras.layers.Dense(1, input_shape = (1, ))
])

class customloss(tf.keras.losses.Loss):
    def call(self, y_true, y_pred):
        return tf.round(tf.square(y_true - y_pred))

model.compile(loss=customloss(),
              optimizer="adam",
              metrics=["mae"])

Raises:

ValueError: The provided loss function (<__main__.customloss object at 0x7f49c5dffeb0>) is not differentiable. 
Training requires a differentiable loss function. Please review your loss function or consider using a standard 
differentiable loss function. You can disable the differentiability check by setting experimental_check_loss_differentiability=False' in 'model.compile()'.

You can see the other usages from this gist.

[haifeng-jin]

will continue the discussion on the issue keras-team/tf-keras#52 to re-scope this PR.

[Frightera]

@haifeng-jin Can you take a look at it again? This now supports custom layers aswell.

While checking custom layers (can throw any error while checking), it uses nested try-except blocks which may not be the best practice.

[Frightera]

Hi @haifeng-jin,

Is there an update for this? I see the questions of new users (new to Keras) tries to use a non-differentiable loss function every two weeks in Stackoverflow.

[haifeng-jin]

Need an review from @qlzh727 since I do not have enough knowledge to review this PR.

[sachinprasadhs]

Hello, Thank you for submitting a pull request.

We're currently in the process of migrating the new Keras 3 code base from keras-team/keras-core to keras-team/keras.
Consequently, merging this PR is not possible at the moment. After the migration is successfully completed, feel free to reopen this PR at keras-team/keras if you believe it remains relevant to the Keras 3 code base. If instead this PR fixes a bug or security issue in legacy tf.keras, you can instead reopen the PR at keras-team/tf-keras, which hosts the TensorFlow-only, legacy version of Keras.