Repository: google-parfait/tensorflow-federated

Issue:
Support for multiple losses with tf.keras
If I understand correctly, tf-federated only supports single loss function. Is there any timeline on when the support for multiple losses will be added?

```
File "/home/.local/lib/python3.5/site-packages/tensorflow_federated/python/learning/model_utils.py", line 456, in __init__
    raise NotImplementedError('only a single loss functions is supported')
```


Discussion:
@aqibsaeed I'll look into how this fits with efforts and follow-up.

Also accepting contributions, if anyone wants to take a stab at doing this :)
@ZacharyGarrett I'd love to take a crack at this. 
@Squadrick thank you, we'd really appreciate any contribution!

Sharing a few notes on things we anticipated may need to be changed (there could definitely be more...):

- [`BatchOutput.loss`](https://github.com/tensorflow/federated/blob/master/tensorflow_federated/python/learning/model.py#L27 ) needs to be updated (at least the pydoc) to explain it is the "total loss" 
- [`tff.learning.from_keras_model`](https://github.com/tensorflow/federated/blob/master/tensorflow_federated/python/learning/keras_utils.py#L87) needs to be updated to allow for a single loss, a `collections.Sequence` of losses, or a `collections.Mapping` of output name → loss function, for the `loss` argument.
- [`_TrainableKerasModel.__init__`](https://github.com/tensorflow/federated/blob/master/tensorflow_federated/python/learning/keras_utils.py#L374) needs to be updated to pass all `inner_model.loss_functions` to the base class
- [`_KerasModel.__init__`](https://github.com/tensorflow/federated/blob/master/tensorflow_federated/python/learning/keras_utils.py#L230) needs a bunch of changes
  - `self._loss_fn` should keep a structure of all loss functions
  - `self._loss_metric` should be updated to match
- [`_KerasModel._forward_pass`](https://github.com/tensorflow/federated/blob/master/tensorflow_federated/python/learning/keras_utils.py#L323) needs to be updated to call each loss function on the associated output, and all losses combined to `batch_loss`

We thought we may want to plumb through a `loss_weights` value from the `from_*_model` methods, similar to Keras' similarly named argument in [`tf.keras.Model.compile`](https://www.tensorflow.org/api_docs/python/tf/keras/models/Model#compile), but perhaps just taking a uniform weighting initially is a good start.

If you're going to be taking a look, I'd recommend a test-driven approach by first adding an example model with two outputs and different losses to [learning/model_examples.py](https://github.com/tensorflow/federated/blob/master/tensorflow_federated/python/learning/model_examples.py#L256), similar to `build_multiple_inputs_keras_model()`, then adding a unittest to [learning/keras_utils_test.py](https://github.com/tensorflow/federated/blob/master/tensorflow_federated/python/learning/keras_utils_test.py#L278) similar to `test_keras_model_multiple_inputs`.
@ZacharyGarrett Thanks for the help, I'll get started.


Resolving patch:
diff --git a/tensorflow_federated/python/learning/keras_utils.py b/tensorflow_federated/python/learning/keras_utils.py
--- a/tensorflow_federated/python/learning/keras_utils.py
+++ b/tensorflow_federated/python/learning/keras_utils.py
@@ -87,6 +87,7 @@ def _preprocess_dummy_batch(dummy_batch):
 def from_keras_model(keras_model,
                      dummy_batch,
                      loss,
+                     loss_weights=None,
                      metrics=None,
                      optimizer=None):
   """Builds a `tff.learning.Model` for an example mini batch.
@@ -98,7 +99,17 @@ def from_keras_model(keras_model,
       The values of the tensors are not important and can be filled with any
       reasonable input value.
     loss: A callable that takes two batched tensor parameters, `y_true` and
-      `y_pred`, and returns the loss.
+      `y_pred`, and returns the loss. If the model has multiple outputs, you
+      can use a different loss on each output by passing a dictionary or a list
+      of losses. The loss value that will be minimized by the model will then
+      be the sum of all individual losses, each weighted by `loss_weights`.
+    loss_weights: (Optional) a list or dictionary specifying scalar coefficients
+      (Python floats) to weight the loss contributions of different model
+      outputs. The loss value that will be minimized by the model will then
+      be the *weighted sum* of all individual losses, weighted by the
+      `loss_weights` coefficients. If a list, it is expected to have a 1:1
+      mapping to the model's outputs. If a tensor, it is expected to map
+      output names (strings) to scalar coefficients.
     metrics: (Optional) a list of `tf.keras.metrics.Metric` objects.
     optimizer: (Optional) a `tf.keras.optimizer.Optimizer`. If None, returned
       model cannot be used for training.
@@ -111,14 +122,51 @@ def from_keras_model(keras_model,
     ValueError: If `keras_model` was compiled.
   """
   py_typecheck.check_type(keras_model, tf.keras.Model)
+  py_typecheck.check_type(
+      loss, (tf.keras.losses.Loss, collections.Sequence, collections.Mapping))
+
+  if loss_weights is not None:
+    py_typecheck.check_type(loss, (collections.Sequence, collections.Mapping))
+
+  if isinstance(loss, (collections.Mapping, collections.Sequence)):
+    if len(loss) != len(keras_model.outputs):
+      raise ValueError('`keras_model` must have equal number of '
+                       'outputs and losses.\nloss: {}\noutputs: {}'.format(
+                           loss, keras_model.outputs))
+    if loss_weights is not None and len(loss) != len(loss_weights):
+      raise ValueError(
+          '`keras_model` must have equal number of '
+          'losses and loss_weights.\nloss: {} \nloss_weights:{}'.format(
+              loss, loss_weights))
+
   if keras_model._is_compiled:  # pylint: disable=protected-access
     raise ValueError('`keras_model` must not be compiled. Use '
                      'from_compiled_keras_model() instead.')
+
   dummy_tensors = _preprocess_dummy_batch(dummy_batch)
   if optimizer is None:
+    if isinstance(loss, collections.Mapping):
+      loss_functions = []
+      for name in keras_model.output_names:
+        if name not in loss:
+          raise KeyError('Output missing from loss dictionary'
+                         '\nlosses: {}\noutputs: {}'.format(
+                             loss.keys(), keras_model.output_names))
+        loss_functions.append(loss[name])
+    elif isinstance(loss, collections.Sequence):
+      loss_functions = loss
+    else:
+      loss_functions = [loss]
+
     return model_utils.enhance(
-        _KerasModel(keras_model, dummy_tensors, loss, metrics))
-  keras_model.compile(loss=loss, optimizer=optimizer, metrics=metrics)
+        _KerasModel(keras_model, dummy_tensors, loss_functions, loss_weights,
+                    metrics))
+
+  keras_model.compile(
+      loss=loss,
+      optimizer=optimizer,
+      loss_weights=loss_weights,
+      metrics=metrics)
   # NOTE: A sub-classed tf.keras.Model does not produce the compiled metrics
   # until the model has been called on input. The work-around is to call
   # Model.test_on_batch() once before asking for metrics.
@@ -230,7 +278,12 @@ def report(accumulators):
 class _KerasModel(model_lib.Model):
   """Internal wrapper class for tf.keras.Model objects."""
 
-  def __init__(self, inner_model, dummy_batch, loss_fn, metrics):
+  def __init__(self,
+               inner_model,
+               dummy_batch,
+               loss_fns,
+               loss_weights=None,
+               metrics=None):
 
     # NOTE: sub-classed `tf.keras.Model`s do not have fully initialized
     # variables until they are called on input. We forced that here.
@@ -246,7 +299,24 @@ def _tensor_spec_with_undefined_batch_dim(tensor):
         _tensor_spec_with_undefined_batch_dim, dummy_batch)
 
     self._keras_model = inner_model
-    self._loss_fn = loss_fn
+    self._loss_fns = loss_fns
+
+    if isinstance(loss_weights, collections.Mapping):
+      self._loss_weights = []
+      for name in inner_model.output_names:
+        if name not in loss_weights:
+          raise KeyError('Output missing from loss_weights dictionary'
+                         '\nloss_weights: {}\noutputs: {}'.format(
+                             loss_weights.keys(), inner_model.output_names))
+        else:
+          self._loss_weights.append(loss_weights[name])
+    else:
+      if loss_weights is None:
+        self._loss_weights = [1.0 for _ in range(len(loss_fns))]
+      else:
+        self._loss_weights = loss_weights
+
+    loss_weights = self._loss_weights
     self._metrics = metrics if metrics is not None else []
 
     # This is defined here so that it closes over the `loss_fn`.
@@ -255,14 +325,24 @@ class _WeightedMeanLossMetric(tf.keras.metrics.Mean):
 
       def __init__(self, name='loss', dtype=tf.float32):
         super(_WeightedMeanLossMetric, self).__init__(name, dtype)
-        self._loss_fn = loss_fn
+        self._loss_fns = loss_fns
+        self._loss_weights = loss_weights
 
       def update_state(self, y_true, y_pred, sample_weight=None):
-        y_true = tf.cast(y_true, self._dtype)
-        y_pred = tf.cast(y_pred, self._dtype)
+        if len(self._loss_fns) == 1:
+          batch_size = tf.cast(tf.shape(y_pred)[0], self._dtype)
+          y_true = tf.cast(y_true, self._dtype)
+          y_pred = tf.cast(y_pred, self._dtype)
+          batch_loss = self._loss_fns[0](y_true, y_pred)
+
+        else:
+          batch_loss = tf.zeros(())
+          for i in range(len(self._loss_fns)):
+            y_t = tf.cast(y_true[i], self._dtype)
+            y_p = tf.cast(y_pred[i], self._dtype)
+            batch_loss += self._loss_weights[i] * self._loss_fns[i](y_t, y_p)
 
-        batch_size = tf.cast(tf.shape(y_pred)[0], self._dtype)
-        batch_loss = self._loss_fn(y_true, y_pred)
+          batch_size = tf.cast(tf.shape(y_pred[0])[0], self._dtype)
 
         return super(_WeightedMeanLossMetric,
                      self).update_state(batch_loss, batch_size)
@@ -331,15 +411,26 @@ def _forward_pass(self, batch_input, training=True):
       raise KeyError('Received a batch_input that is missing required key `x`. '
                      'Instead have keys {}'.format(list(batch_input.keys())))
     predictions = self._keras_model(inputs=inputs, training=training)
-
     y_true = batch_input.get('y')
     if y_true is not None:
-      batch_loss = self._loss_fn(y_true=y_true, y_pred=predictions)
-      for metric in self.get_metrics():
-        metric.update_state(y_true=y_true, y_pred=predictions)
+      if len(self._loss_fns) == 1:
+        loss_fn = self._loss_fns[0]
+        batch_loss = loss_fn(y_true=y_true, y_pred=predictions)
+
+      else:
+        batch_loss = tf.zeros(())
+        for i in range(len(self._loss_fns)):
+          loss_fn = self._loss_fns[i]
+          loss_wt = self._loss_weights[i]
+          batch_loss += loss_wt * loss_fn(
+              y_true=y_true[i], y_pred=predictions[i])
+
     else:
       batch_loss = None
 
+    for metric in self.get_metrics():
+      metric.update_state(y_true=y_true, y_pred=predictions)
+
     return model_lib.BatchOutput(loss=batch_loss, predictions=predictions)
 
   @tf.function
@@ -372,11 +463,9 @@ class _TrainableKerasModel(_KerasModel, model_lib.TrainableModel):
   """Wrapper class for `tf.keras.Model`s that can be trained."""
 
   def __init__(self, inner_model, dummy_batch):
-    if len(inner_model.loss_functions) != 1:
-      raise NotImplementedError('only a single loss functions is supported')
     super(_TrainableKerasModel,
-          self).__init__(inner_model, dummy_batch,
-                         inner_model.loss_functions[0], inner_model.metrics)
+          self).__init__(inner_model, dummy_batch, inner_model.loss_functions,
+                         inner_model.loss_weights, inner_model.metrics)
 
   @property
   def local_variables(self):
diff --git a/tensorflow_federated/python/learning/model.py b/tensorflow_federated/python/learning/model.py
--- a/tensorflow_federated/python/learning/model.py
+++ b/tensorflow_federated/python/learning/model.py
@@ -30,7 +30,8 @@ class BatchOutput(
 
   NOTE: All fields are optional (may be None).
 
-  -   `loss`: The scalar mean loss on the examples in the batch.
+  -   `loss`: The scalar mean loss on the examples in the batch. If the model
+      has multiple losses, it is the sum of all the individual losses.
   -   `predictions`: Tensor of predictions on the examples. The first dimension
       must be the same size (the size of the batch).
   """
diff --git a/tensorflow_federated/python/learning/model_examples.py b/tensorflow_federated/python/learning/model_examples.py
--- a/tensorflow_federated/python/learning/model_examples.py
+++ b/tensorflow_federated/python/learning/model_examples.py
@@ -266,3 +266,18 @@ def build_multiple_inputs_keras_model():
           l.Dense(1)(b),
       ]))
   return tf.keras.Model(inputs=[a, b], outputs=[output])
+
+def build_multiple_outputs_keras_model():
+  l = tf.keras.layers
+  a = l.Input((1, ))
+  b = l.Input((1, ))
+
+  output_a = l.Dense(1)(a)
+  output_b = l.Dense(1)(b)
+  output_c = l.Dense(1)(
+      l.concatenate([
+          l.Dense(1)(a),
+          l.Dense(1)(b)
+      ]))
+
+  return tf.keras.Model(inputs=[a, b], outputs=[output_a, output_b, output_c])


Test patch:
diff --git a/tensorflow_federated/python/learning/keras_utils_test.py b/tensorflow_federated/python/learning/keras_utils_test.py
--- a/tensorflow_federated/python/learning/keras_utils_test.py
+++ b/tensorflow_federated/python/learning/keras_utils_test.py
@@ -451,6 +451,140 @@ def test_keras_model_and_optimizer(self):
     self.assertTrue(hasattr(tff_model._model._keras_model, 'optimizer'))
     # pylint: enable=internal-access
 
+  def test_keras_model_multiple_outputs(self):
+    keras_model = model_examples.build_multiple_outputs_keras_model()
+    dummy_batch = collections.OrderedDict([
+        ('x', [
+            np.zeros([1, 1], dtype=np.float32),
+            np.zeros([1, 1], dtype=np.float32)
+        ]),
+        ('y', [
+            np.zeros([1, 1], dtype=np.float32),
+            np.ones([1, 1], dtype=np.float32),
+            np.ones([1, 1], dtype=np.float32)
+        ]),
+    ])
+
+    with self.subTest("loss_output_len_mismatch"):
+      with self.assertRaises(ValueError):
+        _ = keras_utils.from_keras_model(
+            keras_model=keras_model,
+            dummy_batch=dummy_batch,
+            loss=[
+                tf.keras.losses.MeanSquaredError(),
+                tf.keras.losses.MeanSquaredError()
+            ])
+
+    with self.subTest("invalid_loss"):
+      with self.assertRaises(TypeError):
+        _ = keras_utils.from_keras_model(
+            keras_model=keras_model, dummy_batch=dummy_batch, loss=3)
+
+    with self.subTest("loss_list_no_opt"):
+      tff_model = keras_utils.from_keras_model(
+          keras_model=keras_model,
+          dummy_batch=dummy_batch,
+          loss=[
+              tf.keras.losses.MeanSquaredError(),
+              tf.keras.losses.MeanSquaredError(),
+              tf.keras.losses.MeanSquaredError()
+          ])
+
+      self.assertIsInstance(tff_model, model_utils.EnhancedModel)
+      output = tff_model.forward_pass(dummy_batch)
+      self.assertAllClose(output.loss, 2.0)
+
+    with self.subTest("loss_dict_no_opt"):
+      tff_model = keras_utils.from_keras_model(
+          keras_model=keras_model,
+          dummy_batch=dummy_batch,
+          loss={
+              'dense': tf.keras.losses.MeanSquaredError(),
+              'dense_1': tf.keras.losses.MeanSquaredError(),
+              'dense_2': tf.keras.losses.MeanSquaredError()
+          })
+
+      self.assertIsInstance(tff_model, model_utils.EnhancedModel)
+      output = tff_model.forward_pass(dummy_batch)
+      self.assertAllClose(output.loss, 2.0)
+
+    with self.subTest("trainable_model"):
+      tff_model = keras_utils.from_keras_model(
+          keras_model=keras_model,
+          dummy_batch=dummy_batch,
+          loss=[
+              tf.keras.losses.MeanSquaredError(),
+              tf.keras.losses.MeanSquaredError(),
+              tf.keras.losses.MeanSquaredError()
+          ],
+          optimizer=tf.keras.optimizers.SGD(learning_rate=0.01))
+
+      self.assertIsInstance(tff_model, model_utils.EnhancedTrainableModel)
+      self.assertTrue(hasattr(tff_model._model._keras_model, 'optimizer'))
+      output = tff_model.forward_pass(dummy_batch)
+      self.assertAllClose(output.loss, 2.0)
+
+    keras_model = model_examples.build_multiple_outputs_keras_model()
+    with self.subTest("loss_weights_as_list"):
+      tff_model = keras_utils.from_keras_model(
+          keras_model=keras_model,
+          dummy_batch=dummy_batch,
+          loss=[
+              tf.keras.losses.MeanSquaredError(),
+              tf.keras.losses.MeanSquaredError(),
+              tf.keras.losses.MeanSquaredError()
+          ],
+          loss_weights=[0.1, 0.2, 0.3])
+
+      output = tff_model.forward_pass(dummy_batch)
+      self.assertAllClose(output.loss, 0.5)
+
+    with self.subTest("loss_weights_as_dict"):
+      tff_model = keras_utils.from_keras_model(
+          keras_model=keras_model,
+          dummy_batch=dummy_batch,
+          loss=[
+              tf.keras.losses.MeanSquaredError(),
+              tf.keras.losses.MeanSquaredError(),
+              tf.keras.losses.MeanSquaredError()
+          ],
+          loss_weights={
+              'dense_5': 0.1,
+              'dense_6': 0.2,
+              'dense_7': 0.3
+          })
+
+      output = tff_model.forward_pass(dummy_batch)
+      self.assertAllClose(output.loss, 0.5)
+
+    with self.subTest("loss_weights_assert_fail_list"):
+      with self.assertRaises(ValueError):
+        _ = keras_utils.from_keras_model(
+            keras_model=keras_model,
+            dummy_batch=dummy_batch,
+            loss=[
+                tf.keras.losses.MeanSquaredError(),
+                tf.keras.losses.MeanSquaredError(),
+                tf.keras.losses.MeanSquaredError()
+            ],
+            loss_weights=[0.1, 0.2])
+
+    with self.subTest("loss_weights_assert_fail_dict"):
+      with self.assertRaises(KeyError):
+        _ = keras_utils.from_keras_model(
+            keras_model=keras_model,
+            dummy_batch=dummy_batch,
+            loss=[
+                tf.keras.losses.MeanSquaredError(),
+                tf.keras.losses.MeanSquaredError(),
+                tf.keras.losses.MeanSquaredError()
+            ],
+            loss_weights={
+                'dense_5': 0.1,
+                'dense_6': 0.2,
+                'dummy': 0.4
+            })
+
 
 if __name__ == '__main__':
   test.main()