Pytorch转NCNN的流程记录

带萝卜

发布于 2020-10-26 14:34:31

2.4K0

发布于 2020-10-26 14:34:31

文章被收录于专栏：我的机器学习之路我的机器学习之路我的机器学习之路

最近有一个比较火的ocr项目：chineseocr_lite[1]，项目中很贴心地提供了ncnn的模型推理代码，只需要

交叉编译opencv
添加一点bitmap转cv::Mat的代码
写个简单的界面

具体过程参考：安卓端深度学习模型部署-以NCNN为例 - 带萝卜的文章 - 知乎 https://zhuanlan.zhihu.com/p/137453394

就可以得到一个安卓端的OCR工具了。

可能因为项目针对的是通用的自然场景，所以对小尺寸文本的识别效果不太理想，我对psenet进行了重训练，再转成NCNN进行部署。

PyTorch转NCNN的流程十分简单，如果顺利的话只需要两步：

PyTorch转ONNX

torch.onnx._export(model, x, path, opset_version=11)

ONNX转NCNN

./onnx2ncnn model.onnx model.param model.bin

可是世上哪有那么多一帆风顺的事，这篇文章记录的就是模型重训练之后转成NCNN的过程中遇到的问题和解决方案。

重训练代码参考：wenmuzhou/PSENet.pytorch[2]，backbone选择了mobilenetv3。

注：下文中的resize、interp、interpolate都是一个意思

问题1： ReLU6不支持

概述：ReLU6算子在转换的时候容易出现不支持的情况，需要使用其他算子替代

解决：使用torch.clamp替代（虽然ReLU6可以通过组合ReLU的方式实现，但是组合得到的ReLU6在NCNN中容易转换失败，不建议使用。）

def relu6(x,inplace=True):
     return torch.clamp(x,0,6)

问题2：Resize算子转换问题

概述：因为各个框架对Resize算子的支持都不尽相同[3]，在转换过程中总会出现一些问题，pytorch中的interpolate算子转换成ONNX之后变成很多零散的算子，如cast、shape等，这些在ncnn里面不支持。你可以选择手动修改文件[4]，也可以使用下面这个自动的方法：

解决：使用onnx_simplifier[5]对onnx模型进行简化，可以合并这些零散的算子。

python -m onnxsim model.onnx model_sim.onnx

问题3：关于转ONNX及使用onnx_simplifier过程中出现的一系列奇怪问题

概述：使用不同版本的ONNX可能会遇到不同的问题，比如提示conv层无输入等（具体错误名称记不清了）。

解决：下载最新ONNX源码编译安装（onnx_simplifier中出现的一些错误也可以通过安装最新ONNX来解决）

git clone https://github.com/onnx/onnx.git
sudo apt-get install protobuf-compiler libprotoc-dev 
cd ONNX 
python setup.py install

问题4：模型输出结果的尺寸固定

概述：直接转换得到的onnx模型的Resize算子都是固定输出尺寸的，无论输入多大的图片都会输出同样大小的特征图，这无疑会影响到模型的精度及灵活性。

解决：修改NCNN模型的param文件，将Resize算子修改成按比例resize。

直接转换得到的param文件中的Interp算子是这样的：

Interp    913      1 1 901 913 0=2 1=1.000000e+00 2=1.000000e+00 3=640 4=640

从下面的ncnn源码中可以看到，0代表resize_type，1和2分别是高和宽的缩放比例，3和4分别是输出的高和宽。

int Interp::load_param(const ParamDict& pd) 
{
     resize_type = pd.get(0, 0);
     height_scale = pd.get(1, 1.f);
     width_scale = pd.get(2, 1.f);
     output_height = pd.get(3, 0);
     output_width = pd.get(4, 0);
     return 0; 
}

我们只需将其修改成如下格式即可实现按比例resize：

Interp      913       1 1 901 913 0=1 1=4.000000e+00 2=4.000000e+00

问题5：NCNN模型输出结果与ONNX模型不同

解决：逐层对比NCNN与onnx模型的输出结果

使用onnxruntime（Python）和NCNN（C++）分别提取每个节点的输出，进行对比。对于ncnn比较简单，可以使用

extractor.extract(node_name,preds);

来提取不同节点的输出。

问题5衍生问题：ONNX没有提供提取中间层输出的方法

解决：给要提取的层添加一个输出节点，代码[6]如下：

def find_node_by_name(graph, node_name):
         for node in graph.node:
                 if node.output[0] == node_name:
                         return node
         return None          
def add_extra_output_node(model,target_node, output_name):
         extra_output = helper.make_empty_tensor_value_info(output_name)
         target_output = target_node.output[0]
         identity_node = helper.make_node("Identity",inputs=[target_output],outputs=[output_name],name=output_name)
         model.graph.node.append(identity_node)
         model.graph.output.append(extra_output)
         return model

修改模型之后再使用

out = sess.run([output_name],{"input.1":img.astype(np.float32)})

就可以获取到模型的中间层输出了。

问题5衍生问题：发现最后一个Resize层的输出有差异

解决：参考chineseocr_lite里面的代码把mode由bilinear改成了nearest（这里错误的原因可能是wenmuzhou/PSENet.pytorch中的模型最后一个F.interpolate中的align_corners参数设置成了True。据说NCNN只实现了align_corners为False的情况[7]）。

这里修改之后的模型跟原模型之间是会有少许误差的，如果误差不可接受，就要重新训练才行。