In Tensorflow.keras 2.0, when a model has multiple outputs, how to define a flexible loss function for model.fit()?

Question

Let's say here is a model with two outputs.

import tensorflow as tf
import numpy as np 

x = tf.keras.Input(shape=(35, 7), dtype=tf.float32)     # (None, 35, 7)
net = tf.keras.layers.Dense(11, activation='relu')(x)   # (None, 35, 11)
net = tf.reduce_max(net, axis=1, name='maxpool')        # (None, 11)

a = tf.keras.layers.Dense(13, activation='relu')(net)   # (None, 35, 11)
b = tf.keras.layers.Dense(17, activation='relu')(net)   # (None, 35, 11)
model = tf.keras.Model(inputs=x, outputs=[a, b])

When I do model.compile(loss=loss_fn, optimizer='sgd') :
the model.fit(x=train, y=(label1, label2)) runs loss_fn for each pair of output and label (ie, loss_fn(a, l1) and loss_fn(b, l1) ).

When I do model.compile(loss=[loss_fn1, loss_fn2], optimizer='sgd') :
the model.fit(x=train, y=(label1, label2)) runs loss_fn1 for a and loss_fn2 for b (ie, loss_fn1(a, l1) and loss_fn2(b, l1) ).

So, basically it seems to handle outputs individually (paired with given corresponding labels).

What if I have to define a loss function that should handle/consider both outputs together, and use the function with model.fit ?

(One thing I can think of is to concatenate outputs into one tensor, and separate them in a loss function. However, I don't want to go there since two output may not have consistent shape. Instead, is it possible, for example, something like...)

def loss_fn(y_true, y_pred):
    # I want to access both output ...
    l1, l2 = y_true
    a, b = y_pred
    # ... do something about loss ...
    return loss

Answer 1

You would concatenate your two Dense layers, and do exactly the same as you mentioned:

import numpy as np
from tensorflow.keras.layers import Input, Dense, Concatenate
from tensorflow.keras.models import Model
import tensorflow.keras.backend as K

i = Input((10,))
x = Dense(10)(i)
a = Dense(3, use_bias=False)(x)
b = Dense(3, use_bias=False)(x)
# Now you concatenate both outputs,
# so nothing happens to them
c = Concatenate()([a,b])
m = Model(i, c)

def loss(y_true, y_pred):
    # Do your loss on your subset
    a, b  = y_pred[:, :3], y_pred[:, 3:]
    # Do something random
    return K.mean(a*b)

m.compile("adam", loss)

m.fit(np.random.random((10, 10)),
      np.random.random((10, 6)))

# Outputs:
# 10/10 [=======] - 0s 22ms/sample - loss: -0.2251

edit; haven't seen that actually @bit01 commented already the to go approach