[Python] Keras multioutput custom loss with intermediate layers output

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.

Continue reading...