Keras multioutput custom loss with intermediate layers output

Question

I have a model in keras which takes two inputs and it returns 3 outputs and I want to compute a custom loss. The problem I have is that I don't know how to use the output of intermediate layers in the loss. So far the model consisted of two submodels (submodel1 and submodel2 in the picture) and the final loss consisted of the sum of loss1 and loss2. This was easy because loss1 compared output1 with label1 of the data generator and output2 with label2 of the data generator.

The problem comes when I include the submodel3 in the model, because the loss3 compares the output1 with the output3, being output1 the output of a layer of the model, and not the one that would be the label3 of the data generator. I have tried this way:

input1 = Input(shape=input1_shape)
input2 = Input(shape=input2_shape)
output1 = submodel1()([input1,input2]) #do not pay attention to the code notation, as it is a code to explain the problem.
output2 = submodel2()(output1)
output3 =  submodel3()(output1)
@tf.function
def MyLoss(y_true, y_pred):
    out1, out2, out3 = y_pred
    inp1, inp2 = y_true
            
    loss1 = tf.keras.losses.some_loss1(out1,inp1)
    loss2 = tf.keras.losses.some_loss2(out2, inp2)
    loss3 = tf.keras.losses.some_loss3(out2,out3)

    loss = loss1 + loss2 + loss3
    return loss

model = Model([input1,input2],[output1,output2,output3])
model.compile(optimizer='adam',loss = MyLoss)

But I get this error:

 OperatorNotAllowedInGraphError: iterating over `tf.Tensor` is not allowed: AutoGraph did convert this function. This might indicate you are trying to use an unsupported feature.

I am working with TensorFlow 2.3.0-rc0 version.

Answer 1

You can use add_loss to pass multiple layers output to your custom function. below I replicate your case in a dummy regression task

X1 = np.random.uniform(0,1, (100,5))
X2 = np.random.uniform(0,1, (100,5))

y1 = np.random.uniform(0,1, 100)
y2 = np.random.uniform(0,1, 100)


def MyLoss(true1, true2, out1, out2, out3):

    loss1 = tf.keras.losses.mse(out1, true1)
    loss2 = tf.keras.losses.mse(out2, true2)
    loss3 = tf.keras.losses.mse(out2, out3)

    loss = loss1 + loss2 + loss3
    return loss


input1 = Input(shape=(5,))
input2 = Input(shape=(5,))

output1 = Dense(1)(Concatenate()([input1,input2]))
output2 = Dense(1)(output1)
output3 = Dense(1)(output1)

true1 = Input(shape=(1,))
true2 = Input(shape=(1,))

model = Model([input1,input2,true1,true2], [output1,output2,output3])
model.add_loss(MyLoss(true1, true2, output1, output2, output3))
model.compile(optimizer='adam', loss=None)

model.fit(x=[X1,X2,y1,y2], y=None, epochs=3)

to use the model in inference mode (remove y1, y2 as input):

final_model = Model(model.inputs[:2], model.output)
final_model.predict([X1,X2])