如下图所示,TensorFlow在移动端的上线流程常包括下面几个步骤:
- 固化模型->将variable转constant
- 剪枝->去除、合并、跳过某些节点,
- 量化->数据类型转换
针对预测的优化方向:
- 删除没用的op:
. 只在训练时使用的op
. 不在输入输出路径上的op
. 合并重复节点
. 常量表达式替换- 参数四舍五入
- 预处理BatchNorm
- 转换数据类型
Google的优化工具
- freeze_graph 固化参数,
- summarize_graph 查看模型信息
- optimize_for_inference 剪枝相关的优化(个人理解)
- quantize_graph 量化相关的优化
- 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 负责优化类型,如下
- remove_nodes(op=Identity, op=CheckNumerics) #清除某类无用OP
- fold_constants(ignore_errors=true) #常量表达式替换
- merge_duplicate_nodes #合并重复节点
- strip_unused_nodes(type=float, shape="2,3") #跳过某类无用节点
- fold_batch_norms #预处理BatchNorm
- round_weights(num_steps=256) #参数四舍五入
- quantize_weights #将常量op的float32转int8
- 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/