Using gradients in a custom loss function (tensorflow+keras)

user_3377559 0
Viewed 2118

Is there a natural/easy way in tensorflow/keras to implement a custom loss function that uses the derivatives of the model output with respect to the model input?

I have in mind something along the lines (never mind the actual formula -- this is just a demonstration):

def my_loss_function(y_desired, y_model, x):
  return abs(y_desired - y_model) + abs(tf.gradients(y_model,x[0]))

There are two problems with this. The first is that the loss function doesn't ordinarily have access to the input. My understanding is that I can get around this by directly referencing the Input layer, e.g. as a global variable. (All code is schematic only, of course.)

input_layer = Layer(...)
def my_loss(y1, y2):
  return abs(y1-y2)*input_layer[0]

The second issue is more severe: it seems to not be possible to access the gradient with respect to input_layer, while within the execution graph.

A pretty similar question is here, with no resolution: Custom loss function involving gradients in Keras/Tensorflow. I've made attempts along the same lines, with no luck. (It's not obvious to me that that's the right approach, instead of, say, wrapping Layers in such a way as to always keep track of the derivatives.)

loss-function keras tensorflow
1 Answers

I was not able to implement the training using automatic fit method. However it can certainly be done by manually writing the loop. I will provide example for learning a function only using its gradients.

for epoch in range(epochs):
print("\nStart of epoch %d" % (epoch,))

# Iterate over the batches of the dataset.
for step, (x_batch_train, y_batch_train) in enumerate(train_dataset):

    # Open a GradientTape to record the operations run
    # during the forward pass, which enables auto-differentiation.
    with tf.GradientTape(persistent=True) as tape:

        # Create tensor that you will watch
        x_tensor = tf.convert_to_tensor(x_batch_train, dtype=tf.float64)
        tape.watch(x_tensor)
        # Feed forward
        output = model(x_tensor, training=True)

        # Gradient and the corresponding loss function
        o_x = tape.gradient(output, x_tensor)
        loss_value = loss_fn(y_batch_train, o_x)
    

    # Use the gradient tape to automatically retrieve
    # the gradients of the trainable variables with respect to the loss.
    grads = tape.gradient(loss_value, model.trainable_weights)

    # Run one step of gradient descent by updating
    # the value of the variables to minimize the loss.
    optimizer.apply_gradients(zip(grads, model.trainable_weights))

    # Log every 200 batches.
    if step % 200 == 0:
        print(
            "Training loss (for one batch) at step %d: %.4f"
            % (step, float(loss_value))
        )
        print("Seen so far: %s samples" % ((step + 1) * 64))

loss_fn is simple in this case

loss_fn = tf.keras.losses.MeanSquaredError()

Keep in mind that you are using a second derivative to train the function and the with tf.GradientTape(persistent=True) as tape: will produce a warning which is okay for this case. Also from my experience, this method is particularly sensitive on the choice of an activation function. Continuously differentiable variants of ReLU might be the way to go.

user_25268350
Related