tensorflow: tf.keras cannot weight classes when using multiple outputs

This post is a mirror of https://github.com/keras-team/keras/issues/11735, showing the need to handle class weight for multiple outputs.

Version 2.2.0 used.

---

This is a minimal source code, by @GalAvineri, to reproduce the issue (please comment/uncomment the class weight line):

from tensorflow.python.keras.models import Model
from tensorflow.python.keras.layers import Input, Dense
from tensorflow.python.data import Dataset
import tensorflow as tf
import numpy as np


def preprocess_sample(features, labels):
    label1, label2 = labels
    label1 = tf.one_hot(label1, 2)
    label2 = tf.one_hot(label2, 3)
    return features, (label1, label2)


batch_size = 32

num_samples = 1000
num_features = 10

features = np.random.rand(num_samples, num_features)
labels1 = np.random.randint(2, size=num_samples)
labels2 = np.random.randint(3, size=num_samples)

train = Dataset.from_tensor_slices((features, (labels1, labels2))).map(preprocess_sample).batch(batch_size).repeat()

# Model
inputs = Input(shape=(num_features, ))
output1 = Dense(2, activation='softmax', name='output1')(inputs)
output2 = Dense(3, activation='softmax', name='output2')(inputs)
model = Model(inputs, [output1, output2])

model.compile(loss='categorical_crossentropy', optimizer='adam')
class_weights = {'output1': {0: 1, 1: 10}, 'output2': {0: 5, 1: 1, 2: 10}}
model.fit(train, epochs=10, steps_per_epoch=num_samples // batch_size,
         #  class_weight=class_weights
          )

Uncommenting yields this error:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-38-d137ff6fb3f9> in <module>
     33 class_weights = {'output1': {0: 1, 1: 10}, 'output2': {0: 5, 1: 1, 2: 10}}
     34 model.fit(train, epochs=10, steps_per_epoch=num_samples // batch_size,
---> 35            class_weight=class_weights
     36           )

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py in _method_wrapper(self, *args, **kwargs)
     64   def _method_wrapper(self, *args, **kwargs):
     65     if not self._in_multi_worker_mode():  # pylint: disable=protected-access
---> 66       return method(self, *args, **kwargs)
     67 
     68     # Running inside `run_distribute_coordinator` already.

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py in fit(self, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, sample_weight, initial_epoch, steps_per_epoch, validation_steps, validation_batch_size, validation_freq, max_queue_size, workers, use_multiprocessing)
    813           workers=workers,
    814           use_multiprocessing=use_multiprocessing,
--> 815           model=self)
    816 
    817       # Container that configures and calls `tf.keras.Callback`s.

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/data_adapter.py in __init__(self, x, y, sample_weight, batch_size, steps_per_epoch, initial_epoch, epochs, shuffle, class_weight, max_queue_size, workers, use_multiprocessing, model)
   1115     dataset = self._adapter.get_dataset()
   1116     if class_weight:
-> 1117       dataset = dataset.map(_make_class_weight_map_fn(class_weight))
   1118     self._inferred_steps = self._infer_steps(steps_per_epoch, dataset)
   1119     self._dataset = strategy.experimental_distribute_dataset(dataset)

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/data_adapter.py in _make_class_weight_map_fn(class_weight)
   1233         "Expected `class_weight` to be a dict with keys from 0 to one less "
   1234         "than the number of classes, found {}").format(class_weight)
-> 1235     raise ValueError(error_msg)
   1236 
   1237   class_weight_tensor = ops.convert_to_tensor_v2(

ValueError: Expected `class_weight` to be a dict with keys from 0 to one less than the number of classes, found {'output1': {0: 1, 1: 10}, 'output2': {0: 5, 1: 1, 2: 10}}