From b29f79e052561230f752cd6e61d1292fd976d9a6 Mon Sep 17 00:00:00 2001 From: Tianming Xu Date: Tue, 11 Aug 2020 00:00:53 +0800 Subject: [PATCH] [Relay]Refine tensorflow frontend 1.x & 2.x compatibility (#6240) * [Relay]Refine tensorflow frontend 1.x & 2.x compatibility * fix lint error * revert gpu related changes --- python/tvm/relay/frontend/tensorflow_parser.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/python/tvm/relay/frontend/tensorflow_parser.py b/python/tvm/relay/frontend/tensorflow_parser.py index 771aed0..22c6f94 100644 --- a/python/tvm/relay/frontend/tensorflow_parser.py +++ b/python/tvm/relay/frontend/tensorflow_parser.py @@ -67,12 +67,16 @@ class TFParser(object): def _get_tag_set(self): """Return the tag set of saved model, multiple metagraphs are not supported""" try: - from tensorflow.contrib.saved_model.python.saved_model import reader + from tensorflow.contrib.saved_model.python.saved_model.reader \ + import get_saved_model_tag_sets except ImportError: - raise ImportError( - "InputConfiguration: Unable to import saved_model.reader which is " - "required to get tag set from saved model.") - tag_sets = reader.get_saved_model_tag_sets(self._model_dir) + try: + from tensorflow.python.tools.saved_model_utils import get_saved_model_tag_sets + except ImportError: + raise ImportError( + "InputConfiguration: Unable to import get_saved_model_tag_sets which is " + "required to get tag set from saved model.") + tag_sets = get_saved_model_tag_sets(self._model_dir) return tag_sets[0] def _get_output_names(self): -- 2.7.4