From 4fe8d4a14936dc38558a858283574993909c9895 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Sun, 27 May 2018 10:49:12 -0700 Subject: [PATCH] TPUEstimator.export_savedmodel() saves a SavedModel with both TPU and CPU graphs. PiperOrigin-RevId: 198229550 --- tensorflow/contrib/tpu/python/tpu/tpu_estimator.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py b/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py index 4465833..c8c08a5 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py @@ -1807,7 +1807,7 @@ class TPUEstimator(estimator_lib.Estimator): export_outputs['classes'] = export_output_lib.ClassificationOutput(classes=classes) - tpu.outside_compilation(host_call, logits) + tpu.outside_compilation(host_call, [logits]) ... ``` @@ -1969,7 +1969,7 @@ class TPUEstimator(estimator_lib.Estimator): input_receiver_fn_map[mode]} export_tags = [tag_constants.SERVING, tag_constants.TPU] mode = _REWRITE_FOR_INFERENCE_MODE - try: + if self._export_to_tpu: (super(TPUEstimator, self). _add_meta_graph_for_mode(builder, input_receiver_fn_map, @@ -1978,9 +1978,6 @@ class TPUEstimator(estimator_lib.Estimator): save_variables=False, mode=mode, export_tags=export_tags)) - except Exception as error: # pylint: disable=broad-except - logging.warning('Saving meta graph for TPU failed: {}.' - .format(str(error))) def _call_model_fn(self, features, labels, mode, config): if mode == _REWRITE_FOR_INFERENCE_MODE: -- 2.7.4