使用pnnx把pytorch模型转ncnn模型

序言

在做训练模型转部署模型的时候，通常都是先转成onnx，再转目标框架，但是经常会出现的问题就是某些算子不支持，这样一来，模型部署起来就比较困难，不过现在ncnn支持torch直接通过pnnx转成ncnn模型，将整个torch的模型直接搬过来，避免了算子不支持的问题，具体的介绍还请移步nihui大佬的知乎：PNNX: PyTorch Neural Network Exchange。

结合自己最近的情况，最近在白嫖PaddleOCRV2的移动端模型，想把它转成ncnn去部署在边缘设备上，情况是在转mbv3版本的CRNN的时候推理结果是正常的，但是在转新版本的LCNet版本的CRNN，推理结果错的离谱，因为转换过程异常顺利，所有算子都是支持的，但是转出来的模型识别结果却是差的离谱，使用onnx推理的结果也是正常，我现在都没找到问题的所在，虽然说mbv3版本的crnn识别效果也不错，但是看到paddle官方说LCNet版本的CRNN精度要高好几个点，作为一名有追求的白嫖党，一定要嫖到这个模型！！可能有人会问为什么不用paddle-lite去部署，之前说过了，就自己的项目而言（不代表所有），paddle-lite的部署效率和识别率都远没有ncnn的好，所以这也是为什么执着于ncnn的原因，既然之前的路走不通，那就另寻他路。

说巧不巧，刚好这时候pnnx出来了，我对nihui的崇高仰望已经推到珠穆朗玛峰，太强啦！！！

所以决定试试这一条刚出来的路。本文简单记录一下其中辛酸过程。

一、paddle -> pytorch

前面说了这么多，好像忽略了一个问题，pnnx是pytorch转ncnn，而PaddleOCR是paddle框架的啊，那该咋办？问题不大，paddle和pytorch都是动态图框架，可以相互转换，具体怎么转换简单来说，就是构建相同的网络结构，然后将权重加载进去即可，看似简单，其实坑还是挺大的，知乎上有一篇PaddleOCR识别模型转Pytorch全流程记录示例，其中转换原理可以供参考，新出的LCNet版的CRNN转换也同理。

那么简单说一下我的转换过程吧，首先还是基于PytorchOCR仓库的源码，对网络进行构建，因为这个仓库已经提供了前一个mbv3版本的crnn转换，这样一来也就不用自己去搭轮子从0开始了。使用Netron查看，发现LCNet版的CRNN模型，除了backbone不一样外，head层也由一层全连接，变成了两层全连接层。

考虑到两个框架的api调用方式都差不多，所以直接从paddleocr里将backbone直接复制过来修改，也被了保证构建的层能够保持一致。新建一个RecMv1_enhance.py文件，将paddleocr中的rec_mv1_enhance.py复制过来，将api的调用方式修改成torch的就好。

然后修改heads部分，加入两层全连接，necks不用动，修改完再组合起来就可以得到LCNet作为backbone的crnn模型了，具体怎么组合，希望自己去研究一下这个hub，因为有一个mbv3的例子了，代码理解起来相对来说还是比较容易的。

转换过程中一些细节，在上面给出的文章里有说，我这里比较粗暴，直接将paddle的权重键值修改成torch新构建的网络一致后直接加载（pytorchocr里的mbv3也是这么转的），可以正常识别。

然后将转换后的pytorch权重保存下来。

二、pytorch -> libtorch

pytorch的模型还需要使用jit序列化为libtorch的pt模型后才可以使用pnnx进行转换，这一步比较简单，给个示例程序好了：

def torch2libtorch(model_path,lib_path):
    model = load_model(model_path)[0]
    test_arr = torch.randn(1,3,32,224)
    traced_script_module = torch.jit.trace(model, test_arr)
    output1 = traced_script_module(torch.ones(1, 3, 32, 224))
    output2 = model(torch.ones(1, 3, 32, 224))
    print(output1)
    print(output2)
    traced_script_module.save(lib_path)
    print("->>模型转换成功！")

得到的pt文件，也可以使用libtorch去调用，但是本文的最终目的是转ncnn，所以就先不考虑libtorch调用的问题。

三、pnnx环境准备

pnnx的环境准备参考nihui的知乎博文Windows/Linux/MacOS 编译 PNNX 步骤，支持windows、linux、MacOS，因为我的系统是ubuntu，所以准备步骤如下：

先去pytorch官网下载libtorch文件，然后解压后路径备用，解压路径注意不要有中文：

git clone https://github.com/nihui/ncnn.git
cd ncnn
git checkout pnnx

cd tools/pnnx
mkdir build && cd build

cmake -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=install -DTorch_INSTALL_DIR="<下载位置>/libtorch" ..
cmake --build . -j 2
cmake --build . --target install

编译完后把 libtorch/lib 目录下的 so 全部拷贝到 ncnn/tools/pnnx/build/install/bin 里面，然后将刚才的pt文件拷贝到该目录下，打开终端。运行如下转换命令：

./pnnx MobileNetV1Enhance.pt inputshape=[1,3,32,224] inputshape2=[1,3,32,448]

MobileNetV1Enhance.pt为pt文件路径，两个inputshape的目的是转换出来的模型直接支持动态尺寸输入：

转换顺利的话如上图所示，如果转换过程中遇到一些问题，可以直接在ncnn群里问。转换后会得到几个文件：

除了生成pnnx中间模型之外，还生成了ncnn可以用的param和bin文件，模型调用的时候直接调用ncnn.param和ncnn.bin即可，需要注意输入和输出的名字，默认是in0和out0。

然后转换完后在程序中调用测试一下，在网上随便找了两张假的身份证图片测试，使用onnx2ncnn转换的模型识别效果，可以看到有很多识别错误的符号，我一时半会没找到原因，相同预处理，其他模型转换出来识别都是正常的，唯独这个LCNet版本的CRNN模型转换出来识别效果很差：

然后再看看，使用pnnx转换的模型识别效果，很明显效果要好太多了：

最后再次感慨一下nihui大佬超强的代码工程能力，nihui yyds！！