加载pb文件加载任意模型，并结合tensorboard查看模型结构

本文链接：https://blog.youkuaiyun.com/Fourier_1024/article/details/112317283

前言

之前总结过一篇使用colab训练GitHub中的模型——以tensorflow/models为例。最后生成的为.pb文件。本文将介绍如何使用pb文件。

加载模型

直接上代码。（直接在本地，非colab）

# Load a (frozen) Tensorflow model into memory.
PB_PATH = 'E:/output/pb/DLC/DLC_Model_R0402s.pb'
detection_graph = tf.Graph()
with detection_graph.as_default():
    od_graph_def = tf.GraphDef()
    with tf.gfile.GFile(PB_PATH, 'rb') as fid:
        serialized_graph = fid.read()
        od_graph_def.ParseFromString(serialized_graph)
        tf.import_graph_def(od_graph_def, name='')

使用tensorboard查看模型

tensorboard是tensorflow的配套产品，可以用于图形化查看tensorflow搭建的计算图，也可以监视其运算过程。这里使用它查看计算图的功能。
我使用anaconda安装的tensorboard。
在需要查看计算图时插入：

	writer=tf.summary.FileWriter('E:/output/tensorboard', sess.graph)

在anaconda prompt中输入：

	tensorboard --logdir=E:/output/tensorboard

会弹出一个网站，不要关闭anaconda prompt，将网址复制粘贴到浏览器中查看。
在这里插入图片描述
效果图如下。

根据计算图获取tensor

当得到一个pb文件的时候，例如目标检测模型，我们希望得到一个输入接口，并且获取输出。这些都是可以从计算图中找到的。
以上面的模型为例，其输入为
在这里插入图片描述
输出为：

可以通过get_tensor_by_name来获取相应的接口。代码如下：

        image_tensor = detection_graph.get_tensor_by_name('input_1:0')
        res_tensor = detection_graph.get_tensor_by_name('pred/convolution:0')

之后就可以进行后续的操作了。
上述代码及pb文件已经上传到我的GitHub。