TensorFlow C++ CPU 模型预测优化加速

如下图所示，TensorFlow在移动端的上线流程常包括下面几个步骤：

1. 固化模型->将variable转constant
2. 剪枝->去除、合并、跳过某些节点，
3. 量化->数据类型转换

引用

针对预测的优化方向：

1. 删除没用的op:
   . 只在训练时使用的op
   . 不在输入输出路径上的op
   . 合并重复节点
   . 常量表达式替换
2. 参数四舍五入
3. 预处理BatchNorm
4. 转换数据类型

Google的优化工具

1. freeze_graph 固化参数，
2. summarize_graph 查看模型信息
3. optimize_for_inference 剪枝相关的优化(个人理解)
4. quantize_graph 量化相关的优化
5. graph_transforms 剪枝量化一起上，覆盖3、4功能

工具编译

上述工具均可通过bazel(TensorFlow源码的编译工具)编译，如下

bazel build tensorflow/python/tools:freeze_graph
bazel build tensorflow/tools/graph_transforms:summarize_graph
bazel build tensorflow/python/tools:optimize_for_inference
bazel build tensorflow/tools/quantization:quantize_graph
bazel build tensorflow/tools/graph_transforms:transform_graph

transform_graph工具使用方法

工具1，2的使用都很简单(//todo)，而工具5囊括了工具3、4的功能，且3、4已经被TensorFlow移除，下面重点介绍工具5参数与上述优化方向的对应关系

命令：
bazel-bin/tensorflow/tools/graph_transforms/transform_graph --in_graph=in.pb --out_graph=out.pb --inputs='input1,input2' --outputs='output1' --transforms= ' remove_nodes(op=Identity, op=CheckNumerics) fold_constants(ignore_errors=true) merge_duplicate_nodes strip_unused_nodes(type=float, shape="2,3") fold_batch_norms round_weights(num_steps=256) quantize_weights quantize_nodes'

参数解析：
. in_graph 由工具1 freeze_graph获得
. out_graph 优化后模型输出
. inputs 输入节点，可通过工具2 summarize_graph获得
. outputs 输出节点，可通过工具2 summarize_graph获得
. transforms 负责优化类型，如下

1. remove_nodes(op=Identity, op=CheckNumerics) #清除某类无用OP
2. fold_constants(ignore_errors=true) #常量表达式替换
3. merge_duplicate_nodes #合并重复节点
4. strip_unused_nodes(type=float, shape="2,3") #跳过某类无用节点
5. fold_batch_norms #预处理BatchNorm
6. round_weights(num_steps=256) #参数四舍五入
7. quantize_weights #将常量op的float32转int8
8. quantize_nodes #将所有op的float32转int8

结果

一顿操作猛如虎，结果bert上没啥卵用。最大长度60，CPU使用两个逻辑核，单条预测200ms;CPU使用20个逻辑核时，单条预测80ms。所以最后还是得上资源~~~。后续准备尝试下MKL,XLA等针对编译优化。

原因猜测：
1.剪枝失效->所取tensor即为encoder伸出来的最后一层pooled_ouput，因此没有太大的剪枝空间
2.量化失效->猜测是输入依然为float32，参数即便变成了int8, 运算时依然按照float32进行，因此没啥乱用；也有说现在的量化在x86架构上都不太好使

参考

https://github.com/tensorflow/tensorflow/tree/master/tensorflow/tools/graph_transforms#strip_unused_nodes
https://hanxiao.github.io/2019/01/02/Serving-Google-BERT-in-Production-using-Tensorflow-and-ZeroMQ/