From 138133264c8a75c8e5aed313c27d9e8d85bec4ec Mon Sep 17 00:00:00 2001 From: Zhenyu Tan Date: Tue, 1 Sep 2020 09:06:24 -0700 Subject: [PATCH] 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 | 99 ++++++++++++++++++- .../v1/tensorflow.keras.estimator.pbtxt | 2 +- .../v2/tensorflow.keras.estimator.pbtxt | 2 +- 3 files changed, 97 insertions(+), 6 deletions(-) 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\'], " } }