title: Single loss with Multiple output model in TF.Keras
body: <p>I use <a href="/questions/tagged/tensorflow" class="post-tag" title="show questions tagged &#39;tensorflow&#39;" rel="tag">tensorflow</a>'s Dataset such that <code>y</code> is a dictionary of <strong>6</strong> tensors which I all use in a single loss function which looks likes this:</p>
  <pre class="lang-py prettyprint-override"><code>def CustomLoss():
      def custom_loss(y_true, y_pred):
          a = tf.keras.losses.binary_crossentropy(y_true['a_0'], y_pred[0]) * y_true['a_1']
          b = tf.square(y_true['b_0'] - y_pred[1]) * y_true['b_1']
          c = tf.abs(y_true['c_0'] - y_pred[2]) * y_true['c_1']
          return a + b + c
      return custom_loss
  </code></pre>
  <p>And I have a model with <strong>3</strong> outputs of different shapes. When I compile the model and call <code>fit</code> method I get <code>Value Error</code></p>
  <pre class="lang-py prettyprint-override"><code>model.compile(optimizer=optimizer, loss=CustomLoss())
  model.fit(dataset, epochs=10)
  </code></pre>
  <pre><code>ValueError: Found unexpected keys that do not correspond to any 
  Model output: dict_keys(['a_0', 'a_1', 'b_0', 'b_1', 'c_0', 'c_1']). 
  Expected: ['output_0', 'output_1', 'output_2']
  </code></pre>
  <p>where <code>output_0, 'output_1', 'output_2'</code> are names of the output layers.</p>
  <p>I figured that naming the output layers by the keys in the dataset should solve the issue but the problem is I have <strong>6</strong> tensors in the dataset and only <strong>3</strong> outputs. I'm aware I can assign a loss function to every output with a single dataset ground truth tensor, but again I need to pass at least two tensors as GT.</p>
  <p>So far I've used a custom training loop but I'd rather use the <code>fit</code> method. I'm using <a href="/questions/tagged/tensorflow" class="post-tag" title="show questions tagged &#39;tensorflow&#39;" rel="tag">tensorflow</a> 2.3.1</p>
  <p>EDIT:</p>
  <p>Example model:</p>
  <pre class="lang-py prettyprint-override"><code>inputs = x = tf.keras.layers.Input((256, 256, 3))
  x = tf.keras.applications.ResNet50(include_top=False, weights=None)(x)
  x1 = tf.keras.layers.Flatten()(x)
  x1 = tf.keras.layers.Dense(2, name='output_1')(x1)
  x2 = tf.keras.layers.Conv2D(256, 1, name='output_2')(x)
  x3 = tf.keras.layers.Flatten()(x)
  x3 = tf.keras.layers.Dense(64, name='output_3')(x3)
  model = tf.keras.Model(inputs=inputs, outputs=[x1, x2, x3])
  </code></pre>
  <p>Custom training loop:</p>
  <pre class="lang-py prettyprint-override"><code>avg_loss = tf.keras.metrics.Mean('loss', dtype=tf.float32)
  for epoch in range(1, epochs+1):
      for batch, (images, labels) in enumerate(train_dataset):
          with tf.GradientTape() as tape:
              outputs = model(images, training=False)
              reg_loss = tf.reduce_sum(model.losses)
              pred_loss = loss(labels, outputs)
              total_loss = tf.reduce_sum(pred_loss) + reg_loss
          grads = tape.gradient(total_loss, model.trainable_variables)
          optimizer.apply_gradients(zip(grads, model.trainable_variables))
          avg_loss.update_state(total_loss)
      print(f'Epoch {epoch}/{epochs} - Loss: {avg_loss.result().numpy()}')
      avg_loss.reset_states()
  </code></pre>
  <p>Minimal reproducible code:</p>
  <pre class="lang-py prettyprint-override"><code>import tensorflow as tf
  
  def CustomLoss():
      def custom_loss(y_true, y_pred):
          a = tf.keras.losses.binary_crossentropy(y_true['a_0'], y_pred[0]) * y_true['a_1']
          b = tf.square(y_true['b_0'] - y_pred[1]) * y_true['b_1']
          b = tf.reduce_sum(b, axis=(1, 2, 3))
          c = tf.abs(y_true['c_0'] - y_pred[2]) * y_true['c_1']
          c = tf.reduce_sum(c, axis=1)
          return a + b + c
      return custom_loss
  
  dataset = tf.data.Dataset.from_tensors((
      tf.random.uniform((256, 256, 3)),
      {'a_0': [0., 1.], 'a_1': [1.], 'b_0': tf.random.uniform((8, 8, 256)), 'b_1': [1.], 'c_0': tf.random.uniform((64,)), 'c_1': [1.]}
  ))
  dataset = dataset.batch(1)
  
  inputs = x = tf.keras.layers.Input((256, 256, 3))
  x = tf.keras.applications.ResNet50(include_top=False, weights=None)(x)
  x1 = tf.keras.layers.Flatten()(x)
  x1 = tf.keras.layers.Dense(2, name='output_1')(x1)
  x2 = tf.keras.layers.Conv2D(256, 1, name='output_2')(x)
  x3 = tf.keras.layers.Flatten()(x)
  x3 = tf.keras.layers.Dense(64, name='output_3')(x3)
  model = tf.keras.Model(inputs=inputs, outputs=[x1, x2, x3])
  
  optimizer = tf.keras.optimizers.Adam(1e-3)
  model.compile(optimizer=optimizer, loss=CustomLoss())
  model.fit(dataset, epochs=1)
  </code></pre>