Custom loss function with Gradient Tape, TF2.6-CodePudding

I am trying to use a custom loss function in my Keras sequential model (TensorFlow 2.6.0). This custom loss (ideally) will calculate the data loss plus the residual of a physical equation (say, diffusion equation, Navier Stokes, etc.). This residual error is based on the model output derivative wrt its inputs and I want to use GradientTape.

In this MWE, I removed the data loss term and other equation losses, and just used the derivative of the output wrt its input. The dataset can be found here.

from numpy import loadtxt
from keras.models import Sequential
from keras.layers import Dense
import tensorflow as tf #tf.__version__ = '2.6.0'
# load the dataset
dataset = loadtxt('pima-indians-diabetes.csv', delimiter=',')
# split into input (X) and output (y) variables
X = dataset[:,0:8] #X.shape = (768, 8)
y = dataset[:,8]
X = tf.convert_to_tensor(X, dtype=tf.float32)
y = tf.convert_to_tensor(y, dtype=tf.float32)

def customLoss(y_true,y_pred):
    x_tensor = tf.convert_to_tensor(model.input, dtype=tf.float32)
#     x_tensor = tf.cast(x_tensor, tf.float32)
    with tf.GradientTape() as t:
        t.watch(x_tensor)
        output = model(x_tensor)
    DyDX = t.gradient(output, x_tensor)
    dy_t = DyDX[:, 5:6]
    R_pred=dy_t
    # loss_data = tf.reduce_mean(tf.square(yTrue - yPred), axis=-1)
    loss_PDE = tf.reduce_mean(tf.square(R_pred))
    return loss_PDE

model = Sequential()
model.add(Dense(12, input_dim=8, activation='relu'))
model.add(Dense(12, activation='relu'))
model.add(Dense(12, activation='relu'))
model.add(Dense(1, activation='sigmoid'))

model.compile(loss=customLoss, optimizer='adam', metrics=['accuracy'])

model.fit(X, y, epochs=15)

After execution, I get this ValueError:

ValueError: Passed in object of type <class 'keras.engine.keras_tensor.KerasTensor'>, not tf.Tensor

When I change loss=customLoss to loss='mse', the model starts training, but using that customLoss is the whole point. Any ideas?

CodePudding user response：

The problem seems to come from model.input in the loss function, If I understand your code correctly, you can use the loss :

def custom_loss_pass(model, x_tensor):
    def custom_loss(y_true,y_pred):
        with tf.GradientTape() as t:
            t.watch(x_tensor)
            output = model(x_tensor)
        DyDX = t.gradient(output, x_tensor)
        dy_t = DyDX[:, 5:6]
        R_pred=dy_t
        # loss_data = tf.reduce_mean(tf.square(yTrue - yPred), axis=-1)
        loss_PDE = tf.reduce_mean(tf.square(R_pred))
        return loss_PDE
    return custom_loss

And then:

model.compile(loss=custom_loss_pass(model, X), optimizer='adam', metrics=['accuracy'])

I am not sure it does what you want but at least it works!