1692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt# Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt# 3692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt# Licensed under the Apache License, Version 2.0 (the "License"); 4692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt# you may not use this file except in compliance with the License. 5692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt# You may obtain a copy of the License at 6692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt# 7692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt# http://www.apache.org/licenses/LICENSE-2.0 8692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt# 9692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt# Unless required by applicable law or agreed to in writing, software 10692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt# distributed under the License is distributed on an "AS IS" BASIS, 11692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt# See the License for the specific language governing permissions and 13692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt# limitations under the License. 14692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt# ================================ 15692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt"""Imports a protobuf model as a graph in Tensorboard.""" 16692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt 17692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwaltfrom __future__ import absolute_import 18692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwaltfrom __future__ import division 19692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwaltfrom __future__ import print_function 20692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt 211c42cc48227dab508dae3827f04ec9ef7523bdf4Brian Williammeeimport argparse 221c42cc48227dab508dae3827f04ec9ef7523bdf4Brian Williammeeimport sys 231c42cc48227dab508dae3827f04ec9ef7523bdf4Brian Williammee 24692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwaltfrom tensorflow.core.framework import graph_pb2 25692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwaltfrom tensorflow.python.client import session 26692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwaltfrom tensorflow.python.framework import importer 27692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwaltfrom tensorflow.python.framework import ops 281c42cc48227dab508dae3827f04ec9ef7523bdf4Brian Williammeefrom tensorflow.python.platform import app 29692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwaltfrom tensorflow.python.platform import gfile 30692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwaltfrom tensorflow.python.summary import summary 31692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt 32692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt 33692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwaltdef import_to_tensorboard(model_dir, log_dir): 34692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt """View an imported protobuf model (`.pb` file) as a graph in Tensorboard. 35692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt 36692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt Args: 37692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt model_dir: The location of the protobuf (`pb`) model to visualize 381b5235fd897f7ea5cffc715300f67b4dc852fa27Jonathan Hseu log_dir: The location for the Tensorboard log to begin visualization from. 39692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt 40692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt Usage: 41692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt Call this function with your model location and desired log directory. 42692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt Launch Tensorboard by pointing it to the log directory. 43692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt View your imported `.pb` model as a graph. 44692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt """ 45692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt with session.Session(graph=ops.Graph()) as sess: 46692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt with gfile.FastGFile(model_dir, "rb") as f: 47692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt graph_def = graph_pb2.GraphDef() 48692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt graph_def.ParseFromString(f.read()) 49692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt importer.import_graph_def(graph_def) 50692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt 51692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt pb_visual_writer = summary.FileWriter(log_dir) 52692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt pb_visual_writer.add_graph(sess.graph) 53692fad20f913ffa2cb874a87578ecabb03cc4557Dan Ringwalt print("Model Imported. Visualize by running: " 54b808f1274aab4560ae84a1363acca44fd6e86b54Dan Jarvis "tensorboard --logdir={}".format(log_dir)) 551c42cc48227dab508dae3827f04ec9ef7523bdf4Brian Williammee 561c42cc48227dab508dae3827f04ec9ef7523bdf4Brian Williammee 571c42cc48227dab508dae3827f04ec9ef7523bdf4Brian Williammeedef main(unused_args): 581c42cc48227dab508dae3827f04ec9ef7523bdf4Brian Williammee import_to_tensorboard(FLAGS.model_dir, FLAGS.log_dir) 591c42cc48227dab508dae3827f04ec9ef7523bdf4Brian Williammee 601c42cc48227dab508dae3827f04ec9ef7523bdf4Brian Williammeeif __name__ == "__main__": 611c42cc48227dab508dae3827f04ec9ef7523bdf4Brian Williammee parser = argparse.ArgumentParser() 621c42cc48227dab508dae3827f04ec9ef7523bdf4Brian Williammee parser.register("type", "bool", lambda v: v.lower() == "true") 631c42cc48227dab508dae3827f04ec9ef7523bdf4Brian Williammee parser.add_argument( 641c42cc48227dab508dae3827f04ec9ef7523bdf4Brian Williammee "--model_dir", 651c42cc48227dab508dae3827f04ec9ef7523bdf4Brian Williammee type=str, 661c42cc48227dab508dae3827f04ec9ef7523bdf4Brian Williammee default="", 67c0463c080789854a3cc30e9dbbb27c31f3250b02Dan J required=True, 681c42cc48227dab508dae3827f04ec9ef7523bdf4Brian Williammee help="The location of the protobuf (\'pb\') model to visualize.") 691c42cc48227dab508dae3827f04ec9ef7523bdf4Brian Williammee parser.add_argument( 701c42cc48227dab508dae3827f04ec9ef7523bdf4Brian Williammee "--log_dir", 711c42cc48227dab508dae3827f04ec9ef7523bdf4Brian Williammee type=str, 721c42cc48227dab508dae3827f04ec9ef7523bdf4Brian Williammee default="", 73c0463c080789854a3cc30e9dbbb27c31f3250b02Dan J required=True, 741c42cc48227dab508dae3827f04ec9ef7523bdf4Brian Williammee help="The location for the Tensorboard log to begin visualization from.") 751c42cc48227dab508dae3827f04ec9ef7523bdf4Brian Williammee FLAGS, unparsed = parser.parse_known_args() 761c42cc48227dab508dae3827f04ec9ef7523bdf4Brian Williammee app.run(main=main, argv=[sys.argv[0]] + unparsed) 77