技术学习分享_一航技术 技术资讯 带你了解TensorFlow pb模型常用处理方法

带你了解TensorFlow pb模型常用处理方法

广告位

摘要:TensorFlow 模型训练完成后,通常会通过frozen过程保存一个最终的pb模型。

本文分享自华为云社区《TensorFlow pb模型修改和优化》,作者:luchangli。

TensorFlow 模型训练完成后,通常会通过frozen过程保存一个最终的pb模型。保存的pb模型是以GraphDef数据结构保存的,可以序列化保存为二进制pb模型或者文本pbtxt模型。GraphDef本质上是一个DAG有向无环图,里面主要是存放了一个算子node list,每个算子具有名称,attr等内容,以及通过input包含了node之间的连接关系。

整个GraphDef的输入节点是以Placeholder节点来标识的,模型参数权重通常是以Const节点来保存的。不同于onnx,GraphDef没有对输出进行标识,好处是可以通过node_name:idx来引用获取任意一个节点的输出,缺点是一般需要通过netron手动打开查看模型输出,或者通过代码分析没有输出节点的node作为模型输出节点。下面简单介绍下pb模型常用的一些处理方法。

pb模型保存

# write pb model  with tf.io.gfile.GFile(model_path, "wb") as f:   f.write(graph_def.SerializeToString())  # write pbtxt model  tf.io.write_graph(graph_def, os.path.dirname(model_path), os.path.basename(model_path))

创建node

from tensorflow.core.framework import attr_value_pb2  from tensorflow.core.framework import node_def_pb2  from tensorflow.python.framework import tensor_util  pld_node = node_def_pb2.NodeDef()  pld_node.name = name  pld_node.op = "Placeholder"  shape = tf.TensorShape([None, 3, 256, 256])  pld_node.attr["shape"].CopyFrom(attr_value_pb2.AttrValue(shape=shape.as_proto()))  dtype = tf.dtypes.as_dtype("float32")  pld_node.attr["dtype"].CopyFrom(attr_value_pb2.AttrValue(type=dtype.as_datatype_enum))  # other commonly used setting  node.input.extend(in_node_names)  node.attr["value"].CopyFrom(      attr_value_pb2.AttrValue(tensor=tensor_util.make_tensor_proto(   np_array, np_array.type, np_array.shape)))

构建模型和保存

import tensorflow as tf  import numpy as np  tf.compat.v1.disable_eager_execution()  tf.compat.v1.reset_default_graph()  m = 200  k = 256  n = 128  a_shape = [m, k]  b_shape = [k, n]  np.random.seed(0)  input_np = np.random.uniform(low=0.0, high=1.0, size=a_shape).astype("float32")  kernel_np = np.random.uniform(low=0.0, high=1.0, size=b_shape).astype("float32")  # 构建模型  pld1 = tf.compat.v1.placeholder(dtype="float32", shape=a_shape, name="input1")  kernel = tf.constant(kernel_np, dtype="float32")  feed_dict = {pld1: input_np}  result_tf = tf.raw_ops.MatMul(a=pld1, b=kernel, transpose_a=False, transpose_b=False)  with tf.compat.v1.Session() as sess:      results = sess.run(result_tf, feed_dict=feed_dict)   print("results:", results)  # 保存模型  dump_model_name = "matmul_graph.pb"  graph = tf.compat.v1.get_default_graph()  graph_def = graph.as_graph_def()  with tf.io.gfile.GFile(dump_model_name, "wb") as f:   f.write(graph_def.SerializeToString())

当然一般用其他方式而不是raw_ops构建模型。

pb模型读取

from google.protobuf import text_format  graph_def = tf.compat.v1.GraphDef()  # read pb model  with tf.io.gfile.GFile(model_path, "rb") as f:   graph_def.ParseFromString(f.read())  # read pbtxt model  with open(model_path, "r") as pf:   text_format.Parse(pf.read(), graph_def)

node信息打印

常用信息:

node.name  node.op  node.input  node.device  # please ref https://www.tensorflow.org/api_docs/python/tf/compat/v1/AttrValue  node.attr[attr_name].f # b, i, tensor, etc.  # graph_def中node遍历:  for node in graph_def.node:   ##

对于node的input,一般用node_name:idx如node_name:0来表示输入来自上一个算子的第idx个输出。:0省略则是默认为第0个输出。 名称前面加^符号是控制边。这个input是一个string list,这里面的顺序也对应这个node的各个输入的顺序。

创建GraphDef和添加node

graph_def_n = tf.compat.v1.GraphDef()  for node in graph_def_o.node:   node_n = node_def_pb2.NodeDef()   node_n.CopyFrom(node)   graph_def_n.node.extend([node_n])  # you probably need copy other value like version, etc. from old graph  graph_def_n.version = graph_def_o.version  graph_def_n.library.CopyFrom(graph_def_o.library)  graph_def_n.versions.CopyFrom(graph_def_o.versions)

return graph_def_n

没有onnx模型往graph里面添加节点的topo排序要求

设置placeholder的shape

参考前面创建node部分,通过修改Placeholder的shape属性。

模型shape推导

需要导入模型到tf:tf.import_graph_def(graph_def, name=”)。当然需要先设置正确的pld的shape。

然后获取node的输出tensor:graph.get_tensor_by_name(node_name + “:0”)。

最后可以从tensor里面获取shape和dtype。

pb模型图优化

思路一般比较简单:

1,子图连接关系匹配,比如要匹配conv2d+bn+relu这个pattern连接关系。由于每个node只保存其输入的node连接关系,要进行DFS/BFS遍历图一般需要每个node的输入输出,这可以首先读取所有的node连接关系并根据input信息同时创建一个output信息map。

2,子图替换,先创建新的算子,再把旧的算子替换为新的算子。这个需要创建新的node或者直接修改原来的node。旧的不要的算子可以创建个新图拷贝时丢弃,新的node可以直接extend到graph_def。

3,如果替换为TF内置的算子,算子定义可以参考tensorflow raw_ops中的定义,但是有些属性(例如数据类型attr “T”)没有列出来:https://www.tensorflow.org/api_docs/python/tf/raw_ops

当然也可以替换为自定义算子,这就需要用户开发和注册自定义算子:https://www.tensorflow.org/guide/create_op

如上所述,TensorFlow的pb模型修改优化可以直接使用python代码实现,极大简化开发过程。当然TensorFlow也可以注册grappler和post rewrite图优化pass在C++层面进行图优化,后者除了可以用于推理,也可以用于训练优化。

saved model与pb模型的相互转换

可以参考:tensorflow 模型导出总结 – 知乎

saved model保存的是一整个训练图,并且参数没有冻结。而只用于模型推理serving并不需要完整的训练图,并且参数不冻结无法进行转TensorRT等极致优化。当然也可以saved_model->frozen pb->saved model来同时利用两者的优点。

pb转onnx

使用tf2onnx库GitHub – onnx/tensorflow-onnx: Convert TensorFlow, Keras, Tensorflow.js and Tflite models to ONNX

#!/bin/bash  graphdef=input_model.pb  inputs=Placeholder_1:0,Placeholder_2:0  outputs=output0:0,output1:0  output=${graphdef}.onnx  python -m tf2onnx.convert       --graphdef ${graphdef}       --output ${output}       --inputs ${inputs}       --outputs ${outputs}      --opset 12

 

点击关注,第一时间了解华为云新鲜技术~

本文来自网络,不代表技术学习分享_一航技术立场,转载请注明出处。

作者: 一航技术

上一篇
下一篇
广告位

发表回复

返回顶部