How to implement mid-model loss function with both true and pred values?

Viewed 213

What I'm trying to do is implement the loss functions from this paper.

The loss function utilizes multiple separate models that are linked in series and parallel. The first model is a low-fidelity model, the second model is a set of two, parallel models, and the last model is a physics informed neural network (PINN). My specific implementation will not be using the PINN.

The prescribed loss function needs the MSE of the low-fidelity output and the gradient of the low-fidelity output (that's how I understand the formula at least) as well as the MSE of the high-fidelity output, the MSE of the PINN output, and lastly a regularization term.

So,

  1. How do I get the MSE of a mid-model output?

  2. How do I calculate the gradient terms for y* and y?

  3. Given that there are probably few different ways to implement this, which method would be the best in terms of clarity? (I realize this is a bit subjective)

My own searches have me concluding that I will need to utilize the Model.add_loss() function by either adding it to a subclassed layer or by adding it via the tf.keras.Model API. This all makes sense. The part where I am having issues is getting y_true to the location where I implement the loss function. The best thing I have found so far is the use of dictionaries in model declaration inputs and outputs:

model = tf.keras.Model( inputs = some_dict{'input_name':input_name}, 
            outputs = another_dict{'output_name0': output_name0, 
            'output_name1': output_name1, name = 'my_name' }

I found the general idea here. I did not even know that the use of dicts was even possible for tf.keras.Model inputs prior to reading this. My hang-up is that the author ends up using some of the "outputs" as inputs for evaluation. My current framework looks a bit like this

    def load_model0():
        init_model = [ *Load model of choice* ]
        input0 = tf.keras.Input((size,),name = 'input_zero')
        x = input0
        model0_out = init_model(x)

        model = tf.keras.Model( inputs=input0, outputs = {'model0_out': model0_out}, 
                name = "model0")
        return model

    def load_model1():
        init_model = [ *Load model of choice* ]
        input1 = tf.keras.Input((size,),name = 'input_one')
        x = input0
        model1_out = init_model(x)

        model = tf.keras.Model( inputs=input1, outputs = {'model1_out': model1_out}, 
                name = "model1")
        return model

    def link_models(model0,model1):
        input0 = tf.keras.Input((size0,), name = "first_input")
        input1 = tf.keras.Input((size1,), name = "second_input")

        first_model_out = model0(input0)
        x = tf.keras.layers.Concatenate([first_model_out,input1])
        linked_out = model1(x)

        linked_model = tf.keras.Model( intputs = {'input0': input0, 'input1': input1},
                        outputs = {"linked_out": linked_out, [ *insert confusion* ]}

        loss_function0 = tf.keras.losses.MSE([*insert model0 true values here*],first_model_out)
        linked_model.add_loss(loss_function0)
        loss_function1 = tf.keras.losses.MSE([*insert model1 true values here*], linked_out )
        linked_model.add_loss(loss_function1) 

As of right now, I do not know whether the model.add_loss() loss functions can be used in conjunction with a regular loss function/class in Model.compile(). How do I get the true values? How do I run Model.fit() to make this work?

EDIT: In response to Laplace Ricky. I have a set of three networks connected end to end. The first network doesn't matter here. The second network outputs dimension 15. The third network outputs dimension 40. The setup of the final model looks something like this:

# top/first model takes input dimension of (40,)
inputs = tf.keras.Input((40,))

# middle_out dimension is (size,15)
middle_out = middle_model(top_model(inputs))

linked_model = tf.keras.Model( inputs = inputs,
                   outputs = [bottom_model(middle_out),middle_out],
                   name = 'linked_model' )

I set the loss and optimizer like this:

opt = tf.keras.optimizers.Adam()
loss = [ tf.keras.losses.MSE, tf.keras.losses.MSE ] 
linked_model.compile(optimizer = opt, loss = losses)

and fit with

y_list = [data0,data1]
linked_model.fit(x,y_list, epoch = 10, batch_size = 32 )

where data0 is dimension (size,40), and data1 is dimension (size,15).

This produces the error:

ValueError: Dimensions must be equal, but are 15 and 40 for '{{node Mul_6}} = Mul[T=DT_FLOAT](IteratorGetNext:3, Cast_4)' with input shapes: [160,15], [40].

I have also tried the dict alternative method (with the necessary changes to the expected output in model.fit()):

 loss = {'middle_model': tf.keras.losses.MSE,
         'bottom_model': tf.keras.losses.MSE}
        # Tensorflow gave me a very explicit error saying I need 
        # the name of the bottom model as opposed to the name of the
        # linked model
1 Answers

add_loss() is for losses that don't require the information of the data(y_true) like, for example, regularization losses.

One way to achieve what you want is to create model with multiple outputs.

Example Codes:

tf.random.set_seed(88883)
np.random.seed(88883)
#defining toy example model0 and model1
inputs=tf.keras.Input((10,))
model0=tf.keras.Model(inputs,Dense(5)(inputs),name='outputs_0')
inputs=tf.keras.Input((15,))
model1=tf.keras.Model(inputs,Dense(7)(inputs),name='outputs_1')

#define the big model
inputs_0=tf.keras.Input((10,),name='inputs_0')
inputs_1=tf.keras.Input((10,),name='inputs_1')
first_model_output=model0(inputs_0)
x=tf.keras.layers.Concatenate(axis=-1)([first_model_output,inputs_1])
second_model_output=model1(x)

bigmodel=tf.keras.Model([inputs_0,inputs_1],[first_model_output,second_model_output])

bigmodel.compile(loss=tf.keras.losses.MeanSquaredError(),optimizer=tf.keras.optimizers.SGD())

#generating data
x0=tf.data.Dataset.from_tensor_slices(np.random.rand(64,10))
x1=tf.data.Dataset.from_tensor_slices(np.random.rand(64,10))
y0=tf.data.Dataset.from_tensor_slices(np.random.rand(64,5))
y1=tf.data.Dataset.from_tensor_slices(np.random.rand(64,7))
ds=tf.data.Dataset.zip(({'inputs_0':x0,'inputs_1':x1},{'outputs_0':y0,'outputs_1':y1}))
ds=ds.batch(4)

bigmodel.fit(ds,epochs=5)
'''
Epoch 1/5
16/16 [==============================] - 0s 2ms/step - loss: 0.7648 - outputs_0_loss: 0.4287 - outputs_1_loss: 0.3361
Epoch 2/5
16/16 [==============================] - 0s 1ms/step - loss: 0.5918 - outputs_0_loss: 0.3294 - outputs_1_loss: 0.2625
Epoch 3/5
16/16 [==============================] - 0s 2ms/step - loss: 0.5036 - outputs_0_loss: 0.2728 - outputs_1_loss: 0.2308
Epoch 4/5
16/16 [==============================] - 0s 1ms/step - loss: 0.4530 - outputs_0_loss: 0.2388 - outputs_1_loss: 0.2142
Epoch 5/5
16/16 [==============================] - 0s 2ms/step - loss: 0.4214 - outputs_0_loss: 0.2172 - outputs_1_loss: 0.2042
'''

bigmodel.fit() will be optimizing the sum of MSE of first_model_output with y0 and MSE of second_model_output with y1.

Related