Single loss function with multi-input multi-output model in Keras

Question

I am trying to train a multi-input (3) multi-output (4) model using Keras and I need to use a SINGLE loss function that takes in all the output predictions. 2 of these outputs are my true model outputs that I care about and have corresponding labels, while the other 2 outputs are learnable parameters from within my model that I want to use to dynamically update the loss weights for my true model outputs. I need something like this:

model.compile(optimizer=optimizer, loss = unified_loss

where the unified loss should have access to all my model outputs and corresponding labels. I am using tf.data.from_tensor_slices(...) to train.

The only workaround I have found is to use a custom training loop, which allows this. But, I lose a lot of functionality and callbacks become trickier to implement.

Is there a way to solve this using the regular model.compilt(...) and model.fit(...) ?

Apart from a custom training loop, which is not preferred, I did try the standard approach of:

model.compile(optimizer=optimizer, loss = [loss1, loss2], loss_weights = [alpha, beta]

where I tried to make alpha and beta learnable parameters but this is not desired because I have a custom equation that is more involved than a simple weighted sum.

Answer 1

Add a layer to your model that concats the losses into a single tensor/output. Have your custom loss parse out each of the four values and run the necessary math on them. During inference, run the model without the extra layer.

The pattern of having a slightly different model for training and inference is a common one.