tensorflow(4)使用tensorboard查看ckpt和pb图结构

  在我们使用tensorflow训练和部署模型的时候，我们经常会接触ckpt和pb文件。
  本文将会介绍如何使用tensorboard查看ckpt和pb的模型图结构，便于我们加深对模型的理解。

查看ckpt图结构

  一般tensorflow保存模型格式为ckpt，里面含有data、index、meta文件，其中meta文件为元数据图（meta graph），它保存了tensorflow完整的网络图结构。我们可以通过解析它来查看模型图结构。
  Python代码如下：

import tensorflow as tf
from tensorflow.summary import FileWriter

sess = tf.Session()
tf.train.import_meta_graph("ner.ckpt-8836.meta")
FileWriter("log", sess.graph)
sess.close()

  运行上述代码，会生成log文件夹。接着我们运行如下命令:

tensorboard --logdir=log

在浏览器中输入localhost:6006，即可查看模型图结构，如下图：

查看pb图结构

  谷歌推荐的保存模型的方式是保存模型为pb文件，它具有语言独立性，可独立运行，封闭的序列化格式，任何语言都可以解析它，它允许其他语言和深度学习框架读取、继续训练和迁移TensorFlow的模型。
  首先，我们先使用tensorflow保存一个pb文件，代码如下：

import tensorflow as tf
from tensorflow.python.framework import graph_util


with tf.Session(graph=tf.Graph()) as sess:
    x = tf.placeholder(tf.int32, name='x')
    y = tf.placeholder(tf.int32, name='y')
    b = tf.Variable(1, name='b')
    xy = tf.multiply(x, y)
    # 这里的输出需要加上name属性
    op = tf.add(xy, b, name='op_to_store')

    sess.run(tf.global_variables_initializer())

    # convert_variables_to_constants 需要指定output_node_names，list()，可以多个
    constant_graph = graph_util.convert_variables_to_constants(sess, sess.graph_def, ['op_to_store'])

    # 测试 OP
    feed_dict = {x: 10, y: 3}
    print(sess.run(op, feed_dict))

    # 写入序列化的 PB 文件
    with tf.gfile.FastGFile('test_model.pb', mode='wb') as f:
        f.write(constant_graph.SerializeToString())

  运行上述代码，会生成test_model.pb文件。以下为Python解析pb文件的代码：

from tensorflow.python.platform import gfile
import tensorflow as tf
 
model = 'test_model.pb'
graph = tf.get_default_graph()
graph_def = graph.as_graph_def()
graph_def.ParseFromString(gfile.FastGFile(model, 'rb').read())
tf.import_graph_def(graph_def, name='graph')
summaryWriter = tf.summary.FileWriter('log', graph)

  运行上述代码，会生成log文件夹。接着我们运行如下命令:

tensorboard --logdir=log

在浏览器中输入localhost:6006，即可查看模型图结构，如下图：

  本次分享到此结束。感谢阅读~