Make import_pb_to_tensorboard executable

goat000 · tensorflower-gardener · commit 1c42cc48227d · 2017-06-28T12:04:26.000-07:00
Add main, FLAGS, and BUILD rules.

PiperOrigin-RevId: 160434643
diff --git a/tensorflow/python/tools/BUILD b/tensorflow/python/tools/BUILD
@@ -44,6 +44,20 @@ py_binary(
     ],
 )
 
+py_binary(
+    name = "import_pb_to_tensorboard",
+    srcs = ["import_pb_to_tensorboard.py"],
+    srcs_version = "PY2AND3",
+    deps = [
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python:client",
+        "//tensorflow/python:framework",
+        "//tensorflow/python:framework_ops",
+        "//tensorflow/python:platform",
+        "//tensorflow/python:summary",
+    ],
+)
+
 py_test(
     name = "freeze_graph_test",
     size = "small",
diff --git a/tensorflow/python/tools/import_pb_to_tensorboard.py b/tensorflow/python/tools/import_pb_to_tensorboard.py
@@ -18,10 +18,14 @@
 from __future__ import division
 from __future__ import print_function
 
+import argparse
+import sys
+
 from tensorflow.core.framework import graph_pb2
 from tensorflow.python.client import session
 from tensorflow.python.framework import importer
 from tensorflow.python.framework import ops
+from tensorflow.python.platform import app
 from tensorflow.python.platform import gfile
 from tensorflow.python.summary import summary
 
@@ -48,3 +52,23 @@ def import_to_tensorboard(model_dir, log_dir):
     pb_visual_writer.add_graph(sess.graph)
     print("Model Imported. Visualize by running: "
           "> tensorboard --logdir={}".format(log_dir))
+
+
+def main(unused_args):
+  import_to_tensorboard(FLAGS.model_dir, FLAGS.log_dir)
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser()
+  parser.register("type", "bool", lambda v: v.lower() == "true")
+  parser.add_argument(
+      "--model_dir",
+      type=str,
+      default="",
+      help="The location of the protobuf (\'pb\') model to visualize.")
+  parser.add_argument(
+      "--log_dir",
+      type=str,
+      default="",
+      help="The location for the Tensorboard log to begin visualization from.")
+  FLAGS, unparsed = parser.parse_known_args()
+  app.run(main=main, argv=[sys.argv[0]] + unparsed)
