摘要：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库
#!/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