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