在Linux（Ubantu)下用 PNNX 转换 TorchScript 模型到 ncnn 模型

初次使用 PNNX 的相关步骤记录。

目的

转换 PyTorch 的 TorchScript 格式的模型文件到 ncnn 的模型文件。

环境

• 操作系统： Windows/Ubuntu/macOS 均可，pnnx和pytorch都是跨平台的
• Python 相关:
  • miniconda 安装的 python 版本3.8;
  • PyTorch: 1.11.0
  • torchvision: 0.12.0

# 更新pip，可选
python -m pip install -U pip

# 获取最新版 pytorch， 必须
python -m pip install -U torch

# 如果已经安装了老版本 torchvision， 则需要解决依赖关系：
python -m pip install -U torch torchvision

获取pnnx可执行文件（github下载）

Release Release 20220418 · pnnx/pnnx · GitHub

然后下载Ubantu版本就可以

 使用 PyTorch 导出 TorchScript 格式模型

导出torchscript模型文件

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

比如将save_net.py导出到导出 TorchScript 格式模型

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")

执行

python save_net.py

 得到模型文件

mobilenet_v2.py

用PNNX进行模型转换

将上面Github下的ncnn程序解压到一个文件夹，在mobilenet_v2.py路径下执行

/home/alpha/Desktop/pnnx-20220418-ubuntu/pnnx mobilenet_v2.pt inputshape=[1,3,224,224]
#/home/alpha/Desktop/pnnx-20220418-ubuntu这是上面从Github下载的ncnn程序解压后的路径

得到5个新增文件

mobilenet_v2.ncnn.bin
mobilenet_v2.ncnn.param
mobilenet_v2.pnnx.bin
mobilenet_v2.pnnx.param
mobilenet_v2.py

NCNN模型可视化

netron mobilenet_v2.pnnx.param

 注意 若是netron找不到命令则安装

python -m pip install -U pip
python -m pip install -U netron

运行结果（在Ubantu20.04下）

(base) root@ubuntu:~/Desktop# netron mobilenet_v2.pnnx.param
Serving 'mobilenet_v2.pnnx.param' at http://localhost:8080

打开本地地址

更多参考可能性见原处：

5分钟学会！用 PNNX 转换 TorchScript 模型到 ncnn 模型 - 知乎 (zhihu.com)