• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ================================
15"""Imports a protobuf model as a graph in Tensorboard."""
16
17import argparse
18import sys
19
20from absl import app
21
22from tensorflow.python.client import session
23from tensorflow.python.framework import importer
24from tensorflow.python.framework import ops
25from tensorflow.python.summary import summary
26from tensorflow.python.tools import saved_model_utils
27
28# Try importing TensorRT ops if available
29# TODO(aaroey): ideally we should import everything from contrib, but currently
30# tensorrt module would cause build errors when being imported in
31# tensorflow/contrib/__init__.py. Fix it.
32# pylint: disable=unused-import,g-import-not-at-top,wildcard-import
33try:
34  from tensorflow.contrib.tensorrt.ops.gen_trt_engine_op import *
35except ImportError:
36  pass
37# pylint: enable=unused-import,g-import-not-at-top,wildcard-import
38
39
40def import_to_tensorboard(model_dir, log_dir, tag_set):
41  """View an SavedModel as a graph in Tensorboard.
42
43  Args:
44    model_dir: The directory containing the SavedModel to import.
45    log_dir: The location for the Tensorboard log to begin visualization from.
46    tag_set: Group of tag(s) of the MetaGraphDef to load, in string format,
47      separated by ','. For tag-set contains multiple tags, all tags must be
48      passed in.
49  Usage: Call this function with your SavedModel location and desired log
50    directory. Launch Tensorboard by pointing it to the log directory. View your
51    imported SavedModel as a graph.
52  """
53  with session.Session(graph=ops.Graph()) as sess:
54    input_graph_def = saved_model_utils.get_meta_graph_def(model_dir,
55                                                           tag_set).graph_def
56    importer.import_graph_def(input_graph_def)
57
58    pb_visual_writer = summary.FileWriter(log_dir)
59    pb_visual_writer.add_graph(sess.graph)
60    print("Model Imported. Visualize by running: "
61          "tensorboard --logdir={}".format(log_dir))
62
63
64def main(_):
65  import_to_tensorboard(FLAGS.model_dir, FLAGS.log_dir, FLAGS.tag_set)
66
67
68if __name__ == "__main__":
69  parser = argparse.ArgumentParser()
70  parser.register("type", "bool", lambda v: v.lower() == "true")
71  parser.add_argument(
72      "--model_dir",
73      type=str,
74      default="",
75      required=True,
76      help="The directory containing the SavedModel to import.")
77  parser.add_argument(
78      "--log_dir",
79      type=str,
80      default="",
81      required=True,
82      help="The location for the Tensorboard log to begin visualization from.")
83  parser.add_argument(
84      "--tag_set",
85      type=str,
86      default="serve",
87      required=False,
88      help='tag-set of graph in SavedModel to load, separated by \',\'')
89  FLAGS, unparsed = parser.parse_known_args()
90  app.run(main=main, argv=[sys.argv[0]] + unparsed)
91