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

亦凉 2022-12-19 03:48 280阅读 0赞

在我们使用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，即可查看模型图结构，如下图：  
![查看ckpt模型图结构][ckpt]

### 查看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，即可查看模型图结构，如下图：  
![查看pb图结构][pb]  
  本次分享到此结束。感谢阅读~

[ckpt]: /images/20221120/aac67a82301443ed81981ba32dec31ab.png
[pb]: /images/20221120/03cac761596e44a99aaaa5b9b38006fb.png