[Relay]Refine tensorflow frontend 1.x & 2.x compatibility (#6240)
authorTianming Xu <tianmingxu.tmxu@gmail.com>
Mon, 10 Aug 2020 16:00:53 +0000 (00:00 +0800)
committerGitHub <noreply@github.com>
Mon, 10 Aug 2020 16:00:53 +0000 (09:00 -0700)
* [Relay]Refine tensorflow frontend 1.x & 2.x compatibility

* fix lint error

* revert gpu related changes

python/tvm/relay/frontend/tensorflow_parser.py

index 771aed0..22c6f94 100644 (file)
@@ -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):