摘要：本系列文章旨在分享 tensorflow->onnx->Caffe->wk 模型转换流程，主要针对的是 HI3516CV500， Hi3519AV100 支持 NNIE 推理框架的海思芯片的算法工程落地。
本文分享自华为云社区《将模型转为NNIE框架支持的wk模型——以tensorflow框架为例（一）》，原文作者：wwwyx_*^▽^* 。
使用过 NNIE 框架的同学都知道，NNIE 框架只支持 wk 模型的推理。
实际使用过程中用海思提供的转换软件工具 RuyiStudio 将 caffe 1.0 模型转为 wk。一般情况下如果购买了芯片，海思将会直接将相关的 SDK 包发送给客户，如果没有的话，大家可以从这个链接获取：RuyiStudio
从上面可知，最终需要将别的框架模型转为 caffe 才可以使用 RuyiStudio，目前主流框架包含 pytorch、tensorflow、mxnet 等，关于 pytorch 转 caffe 之前已经有大佬写过了pytorch->caffe，有需要的同学可以去看下。本文主要讲述 tensorflow 框架转为 caffe 可能会遇到的问题以及解决方法。mxnet 有直接的接口可以转 onnx，也可以参考这篇文章做 caffe 转换。
下面进入正题。
tensorflow->caffe
这个真的是个大坑（哭泣），这里我使用了中间模型 onnx，即最终成功转换的路径是 pb->onnx->caffe->wk，下面就说一下具体的操作吧~
第一步：tensorflow->onnx
这一步是最简单的一步= =，目前转了一些模型还没有在这里遇到坑。
使用 github 上的开源项目：tensorflow->onnx，直接使用 pip install 安装后使用。
关注一下都有哪些参数，每个参数的作用，主要是输入、输出、推理使用 nchw 还是 nhwc（caffe 框架为 nchw，所以这里都使用 nchw）、opset(默认使用 9 )，很多的参数我没有使用到，大家有疑问可以直接去 issues 上面看下哈。
下面给出一个转换命令供大家参考下：
python -m tf2onnx.convert --input ./model.pb --inputs input_image:0[1,112,112,3] --inputs-as-nchw input_image:0 --outputs output_0:0,output_1:0,output_2:0,output_3:0,output_4:0 --output ./convert.onnx复制代码
得到 onnx 模型之后，可以使用onnx simplifer将一些零散算子合并，或者将一些冗余算子去除，这个工具视情况使用。
python -m onnxsim input_onnx_model output_onnx_model复制代码
转换为 onnx 之后，需要验证输出的结果是否与 pb 一致，一致后再走后面的流程！！
第二步：onnx->caffe
这里已经得到了 onnx 模型，但是距离成功还有 99%的路要走！！
这一小节 Baseline：onnx2caffe环境:caffe 1.0 + onnx 1.8.0 主要功能代码：
onnx2caffe +-- onnx2caffe | +-- _operators.py | +-- _weightloader.py +-- convertCaffe.py +-- MyCaffe.py复制代码
运行命令：
python convertCaffe.py ./model/MobileNetV2.onnx ./model/MobileNetV2.prototxt ./model/MobileNetV2.caffemodel复制代码
在转换过程中如果遇到了问题，可以从下面几个方面来适配
（1）遇到 caffe 与 NNIE 不支持的算子，可以修改 onnx 模型中的 node 以适配 caffe（这里要发动自己的小脑筋，一些算子替换可以参考一下pytorch->caffe这篇博客）。
（2）如果遇到了 NNIE 与 onnx 支持的算子，但是 caffe 1.0 官方不支持的话，可以在 caffe 中添加新的层，重新编译之后，再做转换。caffe 中添加新的层可以参考：caffe 添加新node
（3）caffe 与 NNIE 都支持的算子，但是转换工具没有支持该算子的转换，在转换代码中添加相应的算子实现。
（4）转换过程中算子转换成功，但是出现了 shape 问题，手动添加一些不需要参数的操作在已经生成的 prototxt 中。
针对上面的每个方法给出对应的解决方式。
修改 onnx 模型中的 node 以适配 caffe
改写 onnx 模型，首先需要了解一下 onnx 都支持哪些算子。
onnx 支持的 op：onnx op
更换模型中的操作时，查看该 node 的输入输出模式，按照格式对模型进行改写。onnx 模型改写涉及多种情况，下面介绍几种常用的方法。
1. 关于 node 的改写有时需要已知其输入输出 size,故一开始先准备一个包含每个 node 输入输出的 onnx 模型。
import onnx.helper as helper
from onnx import shape_inference, TensorProto
import onnxruntime
import onnx
def add_input_output_from_onnx(onnx_path, save_path):
ONNX_DTYPE = {
0: TensorProto.FLOAT,
1: TensorProto.FLOAT,
2: TensorProto.UINT8,
3: TensorProto.INT8,
4: TensorProto.UINT16,
5: TensorProto.INT16,
6: TensorProto.INT32,
7: TensorProto.INT64,
8: TensorProto.STRING,
9: TensorProto.BOOL
}
# load model
onnx_model = onnx.load(onnx_path)
graph = onnx_model.graph
# rewrite the input tensor of graph
input_tensor = graph.input[0]
input_shape = input_tensor.type.tensor_type.shape.dim
input_tensor_new = onnx.helper.make_tensor_value_info(name = input_tensor.name, elem_type = 1,
shape = [1, input_shape[1].dim_value, input_shape[2].dim_value, input_shape[3].dim_value])
graph.input.remove(input_tensor)
graph.input.insert(0, input_tensor_new)
# append all tensor infos to graph input
weight_infos = []
tensors = graph.initializer
for i, tensor in enumerate(tensors):
value_info = helper.make_tensor_value_info(tensor.name, ONNX_DTYPE[tensor.data_type], tensor.dims)
weight_infos.append(value_info)
graph.input.insert(i+1, value_info) # because 0 is for placeholder, so start index is 1
# run node shape inference
node = graph.node
value_info = graph.value_info
inferred_onnx_model = shape_inference.infer_shapes(onnx_model)
onnx.checker.check_model(onnx_model)
inferred_graph = inferred_onnx_model.graph
inferred_value_info = inferred_graph.value_info
onnx.save(inferred_onnx_model,save_path)
return复制代码
使用 netron 打开 onnx 模型，查看添加 size 之后的变化：
2. 遇到 caffe 与 NNIE 不支持的算子，删除 onnx 模型中的 node，将相关操作在外部的预处理阶段进行。这种情况只涉及 onnx 模型中已经存在的节点删除与改变已有边连接的关系，不涉及新的边关系的建立。
` 这里使用graph中node的index来访问node 该代码删除graph node 0,1,2 并且修改node 3的input边 即 input_image --> mul_1 --> sub --> mul --> conv1 变为 input_image --> conv1 ` def delete_node(onnx_path, save_path): onnx_model = onnx.load(onnx_path) graph = onnx_model.graph Mul_1 = graph.node[0] sub = graph.node[1] mul = graph.node[2] conv1 = graph.node[3] conv1.input[0] = Mul_1.input[0] graph.node.remove(Mul_1) graph.node.remove(sub) graph.node.remove(mul) onnx.checker.check_model(onnx_model) onnx.save(onnx_model, save_path)复制代码
3. 更改 caffe 与 NNIE 不支持的算子，修改 onnx 模型中的 node 去适配。如 squeeze 算子，squeeze 算子在 onnx->caffe 的时候会报错，这时可以将 onnx 模型中的 squeeze 替换为 reshape 算子。reshape 需要两个输入，而 squeeze 只对应一个输入，这时需要在 graph 中创建一个新的常数 tensor input。这种情况涉及更换已经存在的 node，新的常数 tensor 的加入，但并不涉及新的边关系的建立。
`查看onnx op的操作，reshape需要两个输入
对于reshape需要将一个shape tensor加入到onnx graph中，
tensor size可以查看第一步生成的onnx model中该squeeze node对应的output size
即 input --> squeeze --> output
变为 input --> reshape(shape) --> output`
def remove_headpose_squeeze_node(onnx_path, save_path):
onnx_model = onnx.load(onnx_path)
graph = onnx_model.graph
## 添加常数 input
shape = onnx.helper.make_tensor('shape', onnx.TensorProto.INT64, [2], [1,3])
graph.initializer.append(shape)
for i in range(len(graph.node)):
if graph.node[i].op_type == "Squeeze":
reshape_node_def = helper.make_node(
'Reshape', # node name
inputs=[graph.node[i].input[0], 'shape'], # inputs
outputs=[graph.node[i].output[0]], # outputs
name = graph.node[i].name
)
graph.node.remove(graph.node[i])
graph.node.insert(i, reshape_node_def)
onnx.checker.check_model(onnx_model)
onnx.save(onnx_model, save_path)复制代码
4. caffe 不支持 div 算子，可以将 div 算子转为 pow+mul。这种情况涉及将一个 node 更换为两个，新的常数 tensor 的加入，以及新的边连接关系。
div 操作： z = x / y
更换为 pow + mul, pow 为幂操作，mul 为乘法操作：
temp = pow(y, -1)
z = temp * x
`
即：
input_x input_y
\\ //
\\ //
div
更改为：
input_x input_y
\\ //
\\ //
\\ pow（常数tensor作为指数输入）
\\ //
\\ // --> （新的边）
mul
`
def change_headpose_div_node(onnx_path, save_path):
onnx_model = onnx.load(onnx_path)
graph = onnx_model.graph
pow_scale = onnx.helper.make_tensor('pow_scale', onnx.TensorProto.FLOAT, [3], [-1.0, -1.0, -1.0])
mul12_output = helper.make_tensor_value_info('pred_pose/mul_12_pow_output:0', onnx.TensorProto.FLOAT, [1, 3])
graph.initializer.append(pow_scale)
# 'pred_pose/mul_12:0' 类似于上图中的input_y
# pow_scale 为上面创建的相应的指数tensor
# 'pred_pose/mul_12_pow_output:0' 为新建的output tensor
# pow name 给一个不与图中node重复的name
mul12_pow_node_def = helper.make_node(
'Pow', # node name
inputs=['pred_pose/mul_12:0', 'pow_scale'], # inputs
outputs=['pred_pose/mul_12_pow_output:0'], # outputs
name = 'pred_pose/mul_12_pow'
)
graph.node.insert(len(graph.node), mul12_pow_node_def)
for i in range(len(graph.node)):
if graph.node[i].name == "pred_pose/truediv_3":
input1 = graph.node[i].input[0]
input2 = graph.node[i].input[1]
output = graph.node[i].output[0]
name = graph.node[i].name
pow_node_def = helper.make_node(
'Mul', # node name
inputs=[input1, mul12_pow_node_def.output[0]], # inputs
outputs=[output], # outputs
name = name
)
print(graph.node[i].name, i)
graph.node.remove(graph.node[i])
graph.node.insert(i, pow_node_def)
break
graph = helper.make_graph(graph.node, graph.name, graph.input, graph.output, graph.initializer)
info_model = helper.make_model(graph)
model = onnx.shape_inference.infer_shapes(info_model)
onnx.save(model, save_path)复制代码
经过这个修改之后，使用 netron 查看 node 边关系，看是否正确。
5. 打印 onnx 中间某个节点的输出，需要在 graph 加一个 output tensor。
def add_outputNode_info(onnx_path, add_name, output_size, save_path): onnx_model = onnx.load(onnx_path) graph = onnx_model.graph prob_info = helper.make_tensor_value_info(add_name,onnx.TensorProto.FLOAT, output_size) graph.output.insert(0, prob_info) onnx.save(onnx_model, save_path) return if __name__ == '__main__': onnx_model = './model.onnx' add_node_path = "./addPreprocessOutput.onnx" # "mul:0": 想要输出node的output name # [1,24,14,14]: 想要输出node的output size add_outputNode_info(onnx_model, "mul:0", [1,24,14,14], add_node_path)复制代码
上面的例子已经将大部分 node 修改的情况涵盖了，修改 onnx 模型可以参考上述代码。
小 tips：Reshape 大法好，各种跟维度有关系的都可以用 reshape 来代替，除此之外，transpose 也是网红 node，具体问题具体分析~
在转换代码中添加相应的算子实现
在 caffe 中添加新的层没什么好说的，按照上面给的链接来就可以，这里主要介绍下如何修改转换代码去适配某个模型转换。经过上面修改 onnx 模型这一步，我们已经将 onnx 模型中的 node 全部换为 caffe 与 NNIE 支持的算子了，但这时 onnx2caffe 可能还会出现问题，下面会从不同的情况做 onnx2caffe 代码适配来逐步完成模型转换。
1. caffe 和 NNIE 都支持某个操作，但是 onnx2caffe 模型转换时报错。
如：TanH 操作，从源码/caffe/src/caffe/layers/中看到有 tanh 层的实现，NNIE 也支持该操作，但是转换报错。查看 onnx2caffe 源码发现没有 TanH 的转换实现，这时需要我们添加相应的转换代码，主要修改_operators.py、_weightloader.py 两个文件，下面以 TanH 为例讲解一下怎么增加转换 node。
_operators.py 文件用来实现 onnx 操作到 Caffe 操作的变换。对于 TanH 的适配，首先需要在文件的最后注册算子模块添加 TanH，然后增加转换代码。
`转换代码：`
def _convert_tanH(node,graph,err):
input_name = str(node.inputs[0])
output_name = str(node.outputs[0])
name = str(node.name)
layer = myf("TanH",name,[input_name],[output_name])
graph.channel_dims[output_name] = graph.channel_dims[input_name]
return layer
`添加注册算子：`
_ONNX_NODE_REGISTRY = {
……
"Tanh": _convert_tanH,
}复制代码
_weightloader.py 文件用来实现 node 参数从 onnx 到 Caffe 的传递。第一步也是在文件末尾添加注册算子，添加同_operators.py。第二步，从 caffe.proto 中查看 tanh 操作是否存在 weight：
message TanHParameter {
enum Engine {
DEFAULT = 0;
CAFFE = 1;
CUDNN = 2;
}
optional Engine engine = 1 [default = DEFAULT];
}复制代码
由于 tanh 操作不存在 weight，所以 onnx 到 caffe 的参数传递为空：
def _convert_tanH(net, node, graph, err): pass复制代码
至此，在 onnx2caffe 中添加 tanh 操作就完成了，具体工程就包含修改上面两个文件夹，主要是注册算子、操作转换的实现、weight 值传递。
caffe 和 NNIE 都支持某个操作，onnx2caffe 也支持该操作，但是操作中有一个输入在模型中被写为 weight，与原来的实现不一致。
如： mul 算子，普通的 mul 算子一般都包含两个输入，模型中可能会存在 mul 算子只有一个输入，另一个输入作为 weight 参数，如下所示：
这种情况下，由于已经存在了 mul 的注册算子，我们只需要在 mul 算子转换的时候新加一个分支来实现就可以了，还是只涉及两个文件的改写。
_operators.py 添加分支代码
def _convert_input1_is_weight_mul(node,graph,max_dim, err):
node_name = node.name
`这里的input_name需要在netron视图中观察一下是哪一个input作为外部输入，这里不能写 weight 的输入名称！`
input_name = str(node.inputs[0])
output_name = str(node.outputs[0])
scale_layer = myf("Scale", node_name, [input_name],[output_name],in_place=False,bias_term=False)
graph.channel_dims[output_name] = max_dim
return scale_layer
def _convert_Mul(node,graph,err):
input_name_list = [str(i) for i in node.inputs]
output_name = str(node.outputs[0])
node_name = node.name
`这里使用node_name 判断mul算子是否是一个input，新增只有一个input的分支`
if node_name == "mul_1":
max_dim = 16
return _convert_input1_is_weight_mul(node,graph,max_dim, err)
···
···复制代码
_weightloader.py 也不需要重新注册，直接添加分支代码
def _convert_input1_is_weight_mul(net, node, graph, err): node_name = node.name ` 注意！！ scale = np.ones(3) * 3.0 对应的是 外部输入size =（1,3）， weight size = (1), 这种情况可以借助 numpy 实现weight与外部输入的channel对齐 这里还有另外一种情况，例如 外部输入 size = (1,128,8,8), weight = (1,128,1,1) 可以这样操作：scale = node.input_tensors[node.inputs[1]] scale = np.reshape(scale, scale.shape[1]) ` scale = np.ones(3) * 3.0 np.copyto(net.params[node_name][0].data, scale, casting='same_kind') `mul本身是没有weight的，所以之前就是直接pass` def _convert_Mul(net, node, graph, err): node_name = node.name if node_name == "mul_1": _convert_input1_is_weight_mul(net, node, graph, err) else： pass复制代码
实际转换过程中，add 算子也会出现上面的情况，其中有一个输入作为算子参数，这时可以把其类比到 _convert_BatchNorm 中的 scale 操作，将 scale 的 weight 视为 1，bias 为 add 算子的内部输入参数，可以参照 BatchNorm 修改代码，这里就不详细写了。
转换过程中算子转换成功，但是出现了 shape 问题,手动修改 prototxt
上面介绍的是算子的适配，但有时通过 onnx2caffe 转换代码之后，已经生成了 prototxt 文件，最终报错 feature map 的 shape 不匹配，由于 onnx2caffe 工具在转换的时候就打印出了每一层的 output，通过与 netron 视图对比，定位第一个出现问题的 node。
知己知彼方能百战百胜，为了定位 shape 为什么不一致，我们先要了解一下不同框架的 padding 策略以及相应的 output size 的计算方法。
- 查看 caffe 的 output size 计算方式，根据代码可得:
output_size=floor((w+2*pad-(d(k-1)+1))/s)+1
template <typename Dtype>
void ConvolutionLayer<Dtype>::compute_output_shape() {
const int* kernel_shape_data = this->kernel_shape_.cpu_data();
const int* stride_data = this->stride_.cpu_data();
const int* pad_data = this->pad_.cpu_data();
const int* dilation_data = this->dilation_.cpu_data();
this->output_shape_.clear();
for (int i = 0; i < this->num_spatial_axes_; ++i) {
// i + 1 to skip channel axis
const int input_dim = this->input_shape(i + 1);
const int kernel_extent = dilation_data[i] * (kernel_shape_data[i] - 1) + 1;
const int output_dim = (input_dim + 2 * pad_data[i] - kernel_extent)/ stride_data[i] + 1;
this->output_shape_.push_back(output_dim);
}
}复制代码
- tensorflow 的 padding 策略可根据这篇博客,结合上面 caffe 的 outputsize 计算，感觉 caffe 的 conv padding 策略与 tensorflow pad=VALID 一致，会把不能参与的 pixel 自动去除不进行计算。
好了，了解了不同框架的 padding 策略以及 output size 的计算方式之后，我们来分析我们的模型，模型转换是这样的：
分析上面模型转换的表格参数：
- tensorflow pad=SAME，为了使所有的 input pixel 都参与计算，tensorflow 在推理时偷偷在 input 的右下补了一行 0，这样最后的输出：
output size = (112 - (1 * (3 - 1) + 1) + 1) / 2 + 1 = 56
其中 (112 - (1 * (3 - 1) + 1)+ 1) 斜体 1 表示偷偷补的 0。
- 对于 onnx，经过查询与实验，发现 pads 参数[0,0,1,1]表示 feature map 上面不补，左边不补，下面补一行 0，右边补一列，与 tf 一致，输出没有什么问题。
- 转为 caffe 之后，caffe 模型 conv pad 参数都为 0，上下左右都不补，这时根据 caffe 的 outputshape 公式，最终计算结果为（1,3,55,55），直接去除 input 的最后一行和最后一列不参与计算。
为了使输出 shape 一致，并且计算结果相同，我采用了下面的解决方法。
caffe 中设置 pad_h:2, pad_w:2。 由于 caffe 是设置 pad 参数之后是对称补 0 的，即 input 的上下左右都补了两行或者两列 0，这时结合 output_shape 公式，最终输出的 shape 为：
output_shape = floor((112 + 2 * 2 - (1 *(3 - 1) + 1) + 1) / 2) + 1 = 57
思考一下 conv 原理，就知道此时 caffe 得到的 featuremap 只是比 tf 的多了最上面一行和最左边一列。稍微解释一下，虽然 caffe 设置 pad=2，但是根据 caffe 的 conv 实现，会将右下比 tf 多补的那一行和那一列自动去除，不参与运算。
这时 feature map 输出为(1,3,57,57), 为了得到正确结果，在 prototxt 文件的 conv 算子之后添加两个 slice 操作，去除最上面一行与最左边一列。
layer {
name: "add_slice1"
type: "Slice"
bottom: "depthwise:0"
top: "add_slice1/split:0"
top: "add_slice1/split:1"
slice_param {
axis: 2
slice_point: 1
}
}
layer {
name: "add_slice2"
type: "Slice"
bottom: "add_slice1/split:1"
top: "add_slice2/split:0"
top: "add_slice2/split:1"
slice_param {
axis: 3
slice_point: 1
}
}复制代码
上面就是针对 caffe 模型的适配，东西很多很杂，有时候需要一些新奇的思路才能解决问题，当然还涉及一些 prototxt 文件中算子 param 的修改，具体问题具体分析，这里就不展开讲了。
第三步：验证
将得到的 caffe 模型的输出结果与 pb 的输出结果进行对比，一般情况下应该是一模一样的，如果不一样主要关注一下输入预处理，输出预处理，被修改的 node 之前的那个 node 的输出是不是 OK（主要是定位是不是自己改的 node 的问题），切忌心浮气躁，掌握方法。每进行一次魔改都做一次推理，这样比较好定位。
总结
对于 tf 转 caffe 确实有一些麻烦，上面可能也只是列了万分之一的问题吧，不过希望可以帮助到大家。大家针对这方面什么好的想法希望可以多交流奥~
针对 onnx 模型的魔改可能是多余的，应该将相关的转换方式直接写进 onnx2caffe 的转换工具中会更加好，但是之前想着修改 onnx 会更简单些，之后希望可以有时间把转换工具修改的更通用一些
强烈要求算法同学训练模型之前先看下 NNIE 框架支持的算子类型！！具体参考《HiSVP 开发指南》5.3.2 节支持的算子类型以及 3.1.6.2 每个算子支持的规格，避免模型转换不过去又要返工！！