从Tensorflow模型文件中解析并显示网络结构图(pb模型篇)

最近看到一个巨牛的人工智能教程,分享一下给大家。教程不仅是零基础,通俗易懂,而且非常风趣幽默,像看小说一样!觉得太牛了,所以分享给大家。平时碎片时间可以当小说看,【点这里可以去膜拜一下大神的“小说”】

Tensorflow官方提供的Tensorboard可以可视化神经网络结构图,但是说实话,我几乎从来不用。主要是因为Tensorboard中查看到的图结构太混乱了,包含了网络中所有的计算节点(读取数据节点、网络节点、loss计算节点等等)。更可怕的是,如果一个计算节点是由多个基础计算(如加减乘除等)构成,那么在Tensorboard中会将基础计算节点显示而不是作为一个整体显示(典型的如Squeeze计算节点)。最近为了排查网络结构BUG花费一周时间,因此,狠下心来决定自己写一个工具,将Tensorflow中的图以最简单的方式显示最关键的网络结构。

1 Tensor对象与Operation对象

Tensorflow中,Tensor对象主要用于存储数据如常量和变量(训练参数),Operation对象是计算节点,如卷积计算、反卷积计算、ReLU等等。每一个Operation对象均有输入和输出Tensor,同理,每个Tensor对象均有对应生成该Tensor的Operation对象和使用该Tensor对象作为输入的Operation对象。Tensor和Operation对象内均有相关属性和函数来获取其关联的Operation和Tensor对象,相关属性如下所示。

Tensor对象的op属性指向生成该Tensor的Operation对象。
Tensor对象的consumers()函数获取使用该Tensor对象作为输入的Operation对象。
Operation对象的inputs属性指向该计算节点的输入Tensor对象。
Operation对象的outputs属性执行该计算节点的输出Tensor对象。

如下图所示的网络结构中,调用Tensor_2对象的consumers()函数,返回的是[op_1,op_2]Tensor_3的op属性指向的是op_1op_1的inputs属性指向的是[Tensor_1,Tensor_2]op_1的output属性指向的是[Tensor_3]

Tensor与Operation

有了Tensor与Operation对应在图中的关联关系,就可以将网络结构给画出来。

2 提取pb文件中的网络结构图

pb文件是将模型参数固化到图文件中,并合并了一些基础计算和删除了反向传播相关计算得到的protobuf协议文件。如果读者还不懂如何将CKPT模型文件转pb文件,请参考我另一篇文章《 Tensorflow MobileNet移植到Android》的第1节部分。有了pb模型文件后,接下来是加载模型,加载pb模型示例代码如下所示。

def read_graph_from_pb(tf_model_path ,input_names,output_name):  
    with open(tf_model_path, 'rb') as f:
        serialized = f.read() 
    tf.reset_default_graph()
    gdef = tf.GraphDef()
    gdef.ParseFromString(serialized) 
    with tf.Graph().as_default() as g:
        tf.import_graph_def(gdef, name='') 
    
    with tf.Session(graph=g) as sess: 
        OPS=get_ops_from_pb(g,input_names,output_name)
    return OPS

其中,倒数第2行调用到的函数get_ops_from_pb()用于获取网络结构图中指定输入节点和指定输出节点之间的计算节点。之所以要指定输入和输出,是为了将输入之前的计算节点(如加载数据队列等相关计算节点)和输出之后的计算节点(如计算loss等相关计算节点)去除,免得碍眼。函数get_ops_from_pb()实现代码如下。

def get_ops_from_pb(graph,input_names,output_name,save_ori_network=True):
    if save_ori_network:
        with open('ori_network.txt','w+') as w: 
            OPS=graph.get_operations()
            for op in OPS:
                txt = str([v.name for v in op.inputs])+'---->'+op.type+'--->'+str([v.name for v in op.outputs])
                w.write(txt+'\n') 
    inputs_tf = [graph.get_tensor_by_name(input_name) for input_name in input_names]
    output_tf =graph.get_tensor_by_name(output_name) 
    OPS =get_ops_from_inputs_outputs(graph, inputs_tf,[output_tf] ) 
    with open('network.txt','w+') as w: 
        for op in OPS:
            txt = str([v.name for v in op.inputs])+'---->'+op.type+'--->'+str([v.name for v in op.outputs])
            w.write(txt+'\n') 
    OPS = sort_ops(OPS)
    OPS = merge_layers(OPS)
    return OPS

在裁剪网络结构(即只保留input_names和output_name之间节点)之前,先将原始的网络结构写入到ori_network.txt中,文件中,每一行写入:输入Tensor---->op---->输出Tensor。接下来调用函数get_ops_from_inputs_outputs获取指定节点之间的节点。并调用sort_ops函数对所有的节点排序,以保证被依赖的节点总是出现在相关节点之前。最后调用merge_layers函数,将一些可以合并的计算合并成一个独立的节点,例如,Squeeze计算相关节点合并成一个单独的Squeeze节点,又如const-->identity两个计算节点可以直接忽略(即删除)。

注意:篇幅有限,这里不再将函数get_ops_from_inputs_outputssort_opsmerge_layers贴出,相关代码请前往文尾提供的源码地址中阅读。

3 绘制网络结构

考虑到SVG绘制图形的简单易用优点,将排好序的网络计算节点和相关Tensor对象数据以Javascript字符串的形式写入到HTML中,使用<line>标签绘制箭头,使用<rect>标签绘制矩形,使用<ellipse>标签绘制椭圆,使用<text>标签显示文字。绘制类似于如下所示图像

绘制网络结构示例

注意:篇幅有限,这里不再介绍Javascript代码解析模型结构和SVG显示相关的原理,相关代码请前往文尾提供的源码地址中阅读。

4 测试模型显示

《MobileNet V1官方预训练模型的使用》文中介绍的MobileNet V1网络结构为例,下载MobileNet_v1_1.0_192文件并压缩后,得到mobilenet_v1_1.0_192_frozen.pb文件。我们还需要知道mobilenet_v1_1.0_192_frozen.pb模型对应的输入和输出Tensor对象的名称,好在MobileNet_v1_1.0_192压缩包中包含文件mobilenet_v1_1.0_192_info.txt。通过该文件可知,输入Tensor的名称为:input:0,输出Tensor名称为:MobilenetV1/Predictions/Reshape_1:0。有了这些信息后,调用函数read_graph_from_pb得到静态图的节点列表对象ops,调用函数gen_graph(ops,"save/path/graph.html")后,在目录save/path中得到graph.html文件,打开graph.html后,显示结果如下。

显示网络结构分两种模式:合并模式和展开模式,分别如下图所示。

合并模式网络结构
截取的展开模式网络结构

5 源码地址

https://github.com/huachao1001/CNNGraph

最后编辑于
©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 158,560评论 4 361
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 67,104评论 1 291
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 108,297评论 0 243
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 43,869评论 0 204
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 52,275评论 3 287
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 40,563评论 1 216
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 31,833评论 2 312
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 30,543评论 0 197
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 34,245评论 1 241
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 30,512评论 2 244
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 32,011评论 1 258
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 28,359评论 2 253
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 33,006评论 3 235
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 26,062评论 0 8
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 26,825评论 0 194
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 35,590评论 2 273
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 35,501评论 2 268

推荐阅读更多精彩内容