diff --git a/tensorflow/python/keras/estimator/__init__.py b/tensorflow/python/keras/estimator/__init__.py
index 895dd0458ef..2f11d1b9f87 100644
--- a/tensorflow/python/keras/estimator/__init__.py
+++ b/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)
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.keras.estimator.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.keras.estimator.pbtxt
index d0dca9a5a31..0a9ee49aecd 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.keras.estimator.pbtxt
+++ b/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\'], "
}
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.keras.estimator.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.keras.estimator.pbtxt
index d9415ba4e54..28d62d03936 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.keras.estimator.pbtxt
+++ b/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\'], "
}
}