Train Neural Network with custom loss-function

Question

i have problems to train my neural network with a custom loss-function. The loss-function i want to use is the following MSE, which consists of MSE_y and MSE_f:

It should be pointed out that the number N_f > N_y. So i want to compute predictions for all train-data and after that i want to calculate my MSE-Function. The values f_i of MSE_f are calculated separately but for simplicity they are just random numbers here (In the code: f). After the calculation of the loss i want to optimize the network. The problem is that i don't know exactly how i can obtain this loss-function. I've tried it to implement it like this:

import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import Dense, Flatten
from tensorflow.keras import Model
from tensorflow.keras.losses import Loss
import matplotlib.pyplot as plt

# Build the tf.keras model using the Keras model subclassing API:
class MyModel(Model):
    def __init__(self):
        super(MyModel, self).__init__()
        self.flatten = Flatten(input_shape=(2, 1))
        self.d1 = Dense(28, activation='sigmoid')
        self.output_ = Dense(1, activation='sigmoid')

    def call(self, x):
        x = self.flatten(x)
        x = self.d1(x)
        x = self.output_(x)
        return x

def myLoss(y_pred, y_true, f):
    loss_func = tf.reduce_mean(tf.square(y_pred-y_true)) + tf.reduce_mean(tf.square(f)) 
    return loss_func


def train(EPOCHS, train_ds, f):
    for epoch in range(EPOCHS):
        # Reset the metrics at the start of the next epoch
        train_loss.reset_states()

        Y_pred = [None] * N_y
        Y_true = [None] * N_y
        i = 0
        with tf.GradientTape() as tape:
            for point, y_true in train_ds:
                y_pred = model(point, training=True)
                Y_pred[i] = y_pred
                Y_true[i] = y_true
                i += 1
            Y_pred = tf.convert_to_tensor(Y_pred, np.float32)
            Y_true = tf.convert_to_tensor(Y_true, np.float32)

            loss = loss_object(Y_true, Y_pred, f)  

        weights = model.trainable_variables
        gradients = tape.gradient(loss, weights)
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))

        train_loss(loss)     
        Loss_history.append(train_loss.result())
        print('Epoch {}, Loss: {}'.format(epoch+1, train_loss.result()))

if __name__ == "__main__":

    np.random.seed(0)
    N_y = 5
    N_f = 10
    # Create N_y= 5 training data samples, each has a x and t-value
    x_train = np.random.rand(N_y, 1, 2, 1).astype("float32")
    y_train = np.random.rand(N_y, 1).astype("float32")

    # Create additional N_f = 10 (for MSE_f)
    x_f_train = np.random.rand(N_f, 1).astype("float32")                       

    #Create tf Datasets
    train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))
    f = tf.convert_to_tensor(x_f_train, np.float32)

    # Create an instance of the model
    model = MyModel()
    optimizer = tf.keras.optimizers.SGD()

    #Loss-Funktion
    loss_object = myLoss

    #I don't know if this metrics is correct for the loss-function?
    train_loss = tf.keras.metrics.Mean(name='train_loss')

    Loss_history = []
    EPOCHS = 10
    train(EPOCHS, train_ds, f)

    plt.figure(1)
    plt.subplot(1, 1, 1)
    plt.plot(Loss_history)
    plt.show()

Is this the right way to train my network with my loss-function MSE? For me it seems that the for-loop and the lists Y_pred, Y_true within the gradientTape are computational not optimal but when i put it outside of the gradientTape there exists no computational graph and therefore the gradients for the optimization are None and nothing works. To put it in a nutshell, how can i optimize my network with my specific loss-function MSE? Thank's for your help :)

---

I use the following configurations:

• Python-Version: 3.7.6
• Tensorflow-Version: 2.1.0
• Keras-Version: 2.2.4-tf