This CL allows model_to_estimator to export signatures for multi-head Keras, or customize signatures for single-head Keras.

PiperOrigin-RevId: 329517928
Change-Id: I89b16614d1fd47b67122026aafb9155dab25cb9e

tensorflow/python/keras/estimator/__init__.py

@@ -38,7 +38,9 @@ def model_to_estimator(
     custom_objects=None,
     model_dir=None,
     config=None,
-    checkpoint_format='saver'):
+    checkpoint_format='saver',
+    metric_names_map=None,
+    export_outputs=None):
   """Constructs an `Estimator` instance from given keras model.
 
   If you use infrastructure or other tooling that relies on Estimators, you can
@@ -71,6 +73,28 @@ def model_to_estimator(
   estimator.train(input_fn, steps=1)
   ```
 
+  Example with customized export signature:
+  ```python
+  inputs = {'a': tf.keras.Input(..., name='a'),
+            'b': tf.keras.Input(..., name='b')}
+  outputs = {'c': tf.keras.layers.Dense(..., name='c')(inputs['a']),
+             'd': tf.keras.layers.Dense(..., name='d')(inputs['b'])}
+  keras_model = tf.keras.Model(inputs, outputs)
+  keras_model.compile(...)
+  export_outputs = {'c': tf.estimator.export.RegressionOutput,
+                    'd': tf.estimator.export.ClassificationOutput}
+
+  estimator = tf.keras.estimator.model_to_estimator(
+      keras_model, export_outputs=export_outputs)
+
+  def input_fn():
+    return dataset_ops.Dataset.from_tensors(
+        ({'features': features, 'sample_weights': sample_weights},
+         targets))
+
+  estimator.train(input_fn, steps=1)
+  ```
+
   Args:
     keras_model: A compiled Keras model object. This argument is mutually
       exclusive with `keras_model_path`. Estimator's `model_fn` uses the
@@ -100,6 +124,32 @@ def model_to_estimator(
       `tf.train.Checkpoint`. Currently, saving object-based checkpoints from
       `model_to_estimator` is only supported by Functional and Sequential
       models. Defaults to 'saver'.
+    metric_names_map: Optional dictionary mapping Keras model output metric
+      names to custom names. This can be used to override the default Keras
+      model output metrics names in a multi IO model use case and provide custom
+      names for the `eval_metric_ops` in Estimator.
+      The Keras model metric names can be obtained using `model.metrics_names`
+      excluding any loss metrics such as total loss and output losses.
+      For example, if your Keras model has two outputs `out_1` and `out_2`,
+      with `mse` loss and `acc` metric, then `model.metrics_names` will be
+      `['loss', 'out_1_loss', 'out_2_loss', 'out_1_acc', 'out_2_acc']`.
+      The model metric names excluding the loss metrics will be
+      `['out_1_acc', 'out_2_acc']`.
+    export_outputs: Optional dictionary. This can be used to override the
+      default Keras model output exports in a multi IO model use case and
+      provide custom names for the `export_outputs` in
+      `tf.estimator.EstimatorSpec`. Default is None, which is equivalent to
+      {'serving_default': `tf.estimator.export.PredictOutput`}. If not None,
+      the keys must match the keys of `model.output_names`.
+      A dict `{name: output}` where:
+        * name: An arbitrary name for this output.
+        * output: an `ExportOutput` class such as `ClassificationOutput`,
+          `RegressionOutput`, or `PredictOutput`. Single-headed models only need
+          to specify one entry in this dictionary. Multi-headed models should
+          specify one entry for each head, one of which must be named using
+          `tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY`
+          If no entry is provided, a default `PredictOutput` mapping to
+          `predictions` will be created.
 
   Returns:
     An Estimator from given keras model.
@@ -126,7 +176,9 @@ def model_to_estimator(
       model_dir=model_dir,
       config=config,
       checkpoint_format=checkpoint_format,
-      use_v2_estimator=False)
+      use_v2_estimator=False,
+      metric_names_map=metric_names_map,
+      export_outputs=export_outputs)
 
 
 @keras_export('keras.estimator.model_to_estimator', v1=[])
@@ -136,7 +188,8 @@ def model_to_estimator_v2(keras_model=None,
                           model_dir=None,
                           config=None,
                           checkpoint_format='checkpoint',
-                          metric_names_map=None):
+                          metric_names_map=None,
+                          export_outputs=None):
   """Constructs an `Estimator` instance from given keras model.
 
   If you use infrastructure or other tooling that relies on Estimators, you can
@@ -169,6 +222,28 @@ def model_to_estimator_v2(keras_model=None,
   estimator.train(input_fn, steps=1)
   ```
 
+  Example with customized export signature:
+  ```python
+  inputs = {'a': tf.keras.Input(..., name='a'),
+            'b': tf.keras.Input(..., name='b')}
+  outputs = {'c': tf.keras.layers.Dense(..., name='c')(inputs['a']),
+             'd': tf.keras.layers.Dense(..., name='d')(inputs['b'])}
+  keras_model = tf.keras.Model(inputs, outputs)
+  keras_model.compile(...)
+  export_outputs = {'c': tf.estimator.export.RegressionOutput,
+                    'd': tf.estimator.export.ClassificationOutput}
+
+  estimator = tf.keras.estimator.model_to_estimator(
+      keras_model, export_outputs=export_outputs)
+
+  def input_fn():
+    return dataset_ops.Dataset.from_tensors(
+        ({'features': features, 'sample_weights': sample_weights},
+         targets))
+
+  estimator.train(input_fn, steps=1)
+  ```
+
   Note: We do not support creating weighted metrics in Keras and converting them
   to weighted metrics in the Estimator API using `model_to_estimator`.
   You will have to create these metrics directly on the estimator spec using the
@@ -248,6 +323,21 @@ def model_to_estimator_v2(keras_model=None,
       `['loss', 'out_1_loss', 'out_2_loss', 'out_1_acc', 'out_2_acc']`.
       The model metric names excluding the loss metrics will be
       `['out_1_acc', 'out_2_acc']`.
+    export_outputs: Optional dictionary. This can be used to override the
+      default Keras model output exports in a multi IO model use case and
+      provide custom names for the `export_outputs` in
+      `tf.estimator.EstimatorSpec`. Default is None, which is equivalent to
+      {'serving_default': `tf.estimator.export.PredictOutput`}. If not None,
+      the keys must match the keys of `model.output_names`.
+      A dict `{name: output}` where:
+        * name: An arbitrary name for this output.
+        * output: an `ExportOutput` class such as `ClassificationOutput`,
+          `RegressionOutput`, or `PredictOutput`. Single-headed models only need
+          to specify one entry in this dictionary. Multi-headed models should
+          specify one entry for each head, one of which must be named using
+          `tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY`
+          If no entry is provided, a default `PredictOutput` mapping to
+          `predictions` will be created.
 
   Returns:
     An Estimator from given keras model.
@@ -275,5 +365,6 @@ def model_to_estimator_v2(keras_model=None,
       config=config,
       checkpoint_format=checkpoint_format,
       use_v2_estimator=True,
-      metric_names_map=metric_names_map)
+      metric_names_map=metric_names_map,
+      export_outputs=export_outputs)
 # LINT.ThenChange(//tensorflow_estimator/python/estimator/keras.py)

tensorflow/tools/api/golden/v1/tensorflow.keras.estimator.pbtxt

@@ -2,6 +2,6 @@ path: "tensorflow.keras.estimator"
 tf_module {
   member_method {
     name: "model_to_estimator"
-    argspec: "args=[\'keras_model\', \'keras_model_path\', \'custom_objects\', \'model_dir\', \'config\', \'checkpoint_format\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\', \'saver\'], "
+    argspec: "args=[\'keras_model\', \'keras_model_path\', \'custom_objects\', \'model_dir\', \'config\', \'checkpoint_format\', \'metric_names_map\', \'export_outputs\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\', \'saver\', \'None\', \'None\'], "
   }
 }

tensorflow/tools/api/golden/v2/tensorflow.keras.estimator.pbtxt

@@ -2,6 +2,6 @@ path: "tensorflow.keras.estimator"
 tf_module {
   member_method {
     name: "model_to_estimator"
-    argspec: "args=[\'keras_model\', \'keras_model_path\', \'custom_objects\', \'model_dir\', \'config\', \'checkpoint_format\', \'metric_names_map\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\', \'checkpoint\', \'None\'], "
+    argspec: "args=[\'keras_model\', \'keras_model_path\', \'custom_objects\', \'model_dir\', \'config\', \'checkpoint_format\', \'metric_names_map\', \'export_outputs\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\', \'checkpoint\', \'None\', \'None\'], "
   }
 }