使用pnnx将Torch模型转换为ncnn

1. 引言

以往我们将Torch模型转换为ncnn模型，通常需经过Torch–>onnx，onnx–>ncnn两个过程。但经常会出现某些算子不支持的问题。
ncnn作者针对该问题，直接开发一个Torch直接转换ncnn模型的工具 (PNNX)，以下为相关介绍及使用方法。

2. 介绍

PyTorch神经网络eXchange（PNNX）是PyTorch模型互操作性的开放标准。PNNX为PyTorch提供了一种开放的模型格式。它定义了计算图以及严格匹配PyTorch的高级运算符。

3. 使用

3.1 安装

• Windows/Linux/macOS 64bit
• python 3.7 or later

pip3 install pnnx

从下述链接下载pnnx预编译包：
https://github.com/pnnx/pnnx/releases

此包包括所需的所有二进制文件。它是可移植的，因此不需要CUDA或PyTorch运行时环境

3.2 使用 PyTorch 导出 TorchScript 格式模型

导出torchscript模型文件

第一步是从 PyTorch 导出 TorchScript 模型， 例如 mobilenet_v2 这一经典的网络:
save_net.py:

import torch
import torchvision.models as models

#net = models.resnet18(pretrained=True)
net = models.mobilenet_v2(pretrained=True)
net = net.eval()

x = torch.rand(1, 3, 224, 224)

mod = torch.jit.trace(net, x)
torch.jit.save(mod, "mobilenet_v2.pt")

执行可得到模型文件 mobilenet_v2.pt。

3.3 用PNNX把torchscript模型转为ncnn模型

./pnnx mobilenet_v2.pt inputshape=[1,3,224,224]

运行后会生成下列文件

• mobilenet_v2.pnnx.param PNNX graph definition
• mobilenet_v2.pnnx.bin PNNX model weight
• mobilenet_v2_pnnx.py PyTorch script for inference, the python code - for model construction and weight initialization
• mobilenet_v2.pnnx.onnx PNNX model in onnx format
• mobilenet_v2.ncnn.param ncnn graph definition
• mobilenet_v2.ncnn.bin ncnn model weight
• mobilenet_v2_ncnn.py pyncnn script for inference

3.4 模型可视化

直接用 netron 查看即可

4. 参考

https://blog.csdn.net/qq_39056987/article/details/121331430
https://zhuanlan.zhihu.com/p/427620428
https://zhuanlan.zhihu.com/p/427512763
https://github.com/pnnx/pnnx