how to make a custom loss function which use model in keras

Question

I'm trying to make a custom loss function for keras NN model. Normally, loss functions have y_prediction and y_true for arguments. But, I need to use model in the custom loss function like y_prediction = model(X_train) to use tf. GradientTape tf. GradientTape . So what I want to know is how to use the latest model(on the way to fit) in the custom loss function.

If you have an idea about that, tell me, please. (Sorry for my bad English)

Answer 1

You can create a model class as and implement the train_step method:

class YourModel(Model):
    def __init__(self):
        super(YourModel, self).__init__()

        # define your model architecture here as an attribute of the class

  def train_step(data):
      with tf.GradientTape() as tape:
          # foward pass data through the architecture
          # compute loss (y_true, y_pred, any other param)
      
      # weight update
      gradients = tape.gradient(loss, self.trainable_variables) 
      self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))

      return {
          'loss': loss
          # other losses
      }

  def call(self, x):
    # your forward pass implementation

    return # output

More information can be found here: https://www.tensorflow.org/tutorials/quickstart/advanced