带你了解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

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