Skip to content

Commit

Permalink
[Relay]Refine tensorflow frontend 1.x & 2.x compatibility (apache#6240)
Browse files Browse the repository at this point in the history
* [Relay]Refine tensorflow frontend 1.x & 2.x compatibility

* fix lint error

* revert gpu related changes
  • Loading branch information
xutianming authored and trevor-m committed Sep 3, 2020
1 parent 380e561 commit bfc1d97
Showing 1 changed file with 9 additions and 5 deletions.
14 changes: 9 additions & 5 deletions python/tvm/relay/frontend/tensorflow_parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -67,12 +67,16 @@ def _load_pb_file(self):
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):
Expand Down

0 comments on commit bfc1d97

Please sign in to comment.