Pytorch转onnx、torchscript方式


Posted in Python onMay 25, 2020

前言

本文将介绍如何使用ONNX将PyTorch中训练好的模型(.pt、.pth)型转换为ONNX格式,然后将其加载到Caffe2中。需要安装好onnx和Caffe2。

PyTorch及ONNX环境准备

为了正常运行ONNX,我们需要安装最新的Pytorch,你可以选择源码安装:

git clone --recursive https://github.com/pytorch/pytorch
cd pytorch
mkdir build && cd build
sudo cmake .. -DPYTHON_INCLUDE_DIR=/usr/include/python3.6 -DUSE_MPI=OFF
make install
export PYTHONPATH=$PYTHONPATH:/opt/pytorch/build

其中 "/opt/pytorch/build"是前面build pytorch的目。

or conda安装

conda install pytorch torchvision -c pytorch

安装ONNX的库

conda install -c conda-forge onnx

onnx-caffe2 安装

pip3 install onnx-caffe2

Pytorch模型转onnx

在PyTorch中导出模型通过跟踪工作。要导出模型,请调用torch.onnx.export()函数。这将执行模型,记录运算符用于计算输出的轨迹。因为_export运行模型,我们需要提供输入张量x。

这个张量的值并不重要; 它可以是图像或随机张量,只要它是正确的大小。更多详细信息,请查看torch.onnx documentation文档。

# 输入模型
example = torch.randn(batch_size, 1, 224, 224, requires_grad=True)

# 导出模型
torch_out = torch_out = torch.onnx.export(model, # model being run
    example, # model input (or a tuple for multiple inputs)
    "peleeNet.onnx",
 verbose=False, # store the trained parameter weights inside the model file
 training=False,
 do_constant_folding=True,
 input_names=['input'],
 output_names=['output'])

其中torch_out是执行模型后的输出,通常以忽略此输出。转换得到onnx后可以使用OpenCV的 cv::dnn::readNetFromONNX or cv::dnn::readNet进行模型加载推理了。

还可以进一步将onnx模型转换为ncnn进而部署到移动端。这就需要ncnn的onnx2ncnn工具了.

编译ncnn源码,生成 onnx2ncnn。

其中onnx转换模型时有一些冗余,可以使用用工具简化一些onnx模型。

pip3 install onnx-simplifier

简化onnx模型

python3 -m onnxsim pnet.onnx pnet-sim.onnx

转换成ncnn

onnx2ncnn pnet-sim.onnx pnet.param pnet.bin

ncnn 加载模型做推理

Pytorch模型转torch script

pytorch 加入libtorch前端处理,集体步骤为:

Pytorch转onnx、torchscript方式

以mtcnn pnet为例

# convert pytorch model to torch script
# An example input you would normally provide to your model's forward() method.
example = torch.rand(1, 3, 12, 12).to(device)
# Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing.
traced_script_module = torch.jit.trace(pnet, example)
# Save traced model
traced_script_module.save("pnet_model_final.pth")

C++调用如下所示:

#include <torch/script.h> // One-stop header.
#include <iostream>
#include <memory>
int main(int argc, const char* argv[]) 
{
 if (argc != 2) 
 {
 std::cerr << "usage: example-app <path-to-exported-script-module>\n";
 return -1;
 }

 // Deserialize the ScriptModule from a file using torch::jit::load().
 std::shared_ptr<torch::jit::script::Module> module = torch::jit::load(argv[1]);

 assert(module != nullptr);
 std::cout << "ok\n";
}
Python 相关文章推荐
Python中的fileinput模块的简单实用示例
Jul 09 Python
python实现下载指定网址所有图片的方法
Aug 08 Python
Python线程指南详细介绍
Jan 05 Python
零基础使用Python读写处理Excel表格的方法
May 02 Python
Puppeteer使用示例详解
Jun 20 Python
python中的数组赋值与拷贝的区别详解
Nov 26 Python
django xadmin 管理器常用显示设置方式
Mar 11 Python
在django admin中配置搜索域是一个外键时的处理方法
May 20 Python
详解Django ORM引发的数据库N+1性能问题
Oct 12 Python
Python 无限级分类树状结构生成算法的实现
Jan 21 Python
Python利用socket模块开发简单的端口扫描工具的实现
Jan 27 Python
Python实现生活常识解答机器人
Jun 28 Python
使用pandas库对csv文件进行筛选保存
May 25 #Python
pytorch中 gpu与gpu、gpu与cpu 在load时相互转化操作
May 25 #Python
基于pandas向csv添加新的行和列
May 25 #Python
Python如何把十进制数转换成ip地址
May 25 #Python
tensorflow模型转ncnn的操作方式
May 25 #Python
MxNet预训练模型到Pytorch模型的转换方式
May 25 #Python
浅谈pytorch 模型 .pt, .pth, .pkl的区别及模型保存方式
May 25 #Python
You might like
几种显示数据的方法的比较
2006/10/09 PHP
WordPress判断用户是否登录的代码
2011/03/17 PHP
PHP单例模式实例分析【防继承,防克隆操作】
2019/05/22 PHP
基于jquery DOM写的类似微博发布的效果
2012/10/20 Javascript
密码强度检测效果实现原理与代码
2013/01/04 Javascript
js改变文章字体大小的实例代码
2013/11/27 Javascript
JSONP获取Twitter和Facebook文章数的具体步骤
2014/02/24 Javascript
JavaScript中的console.log()函数详细介绍
2014/12/29 Javascript
javascript自定义右键弹出菜单实现方法
2015/05/25 Javascript
vue.js+boostrap项目实践(案例详解)
2016/09/21 Javascript
JavaScript动态数量的文件上传控件
2016/11/18 Javascript
利用js获取下拉框中所选的值
2016/12/01 Javascript
es6学习笔记之Async函数基本教程
2017/05/11 Javascript
基于Vue的文字跑马灯组件(npm 组件包)
2017/05/24 Javascript
基于Vuejs和Element的注册插件的编写方法
2017/07/03 Javascript
jQuery 防止相同的事件快速重复触发方法
2018/02/08 jQuery
[06:43]2018DOTA2国际邀请赛寻真——VGJ.Thunder
2018/08/11 DOTA
[51:17]VGJ.T vs Mineski 2018国际邀请赛小组赛BO2 第二场 8.18
2018/08/19 DOTA
用Python编程实现语音控制电脑
2014/04/01 Python
Linux 下 Python 实现按任意键退出的实现方法
2016/09/25 Python
pycharm安装图文教程
2017/05/02 Python
转换科学计数法的数值字符串为decimal类型的方法
2018/07/16 Python
详解python中list的使用
2019/03/15 Python
Python爬取数据保存为Json格式的代码示例
2019/04/09 Python
pytorch在fintune时将sequential中的层输出方法,以vgg为例
2019/08/20 Python
Python爬虫运用正则表达式的方法和优缺点
2019/08/25 Python
pytorch中获取模型input/output shape实例
2019/12/30 Python
Python callable内置函数原理解析
2020/03/05 Python
python topk()函数求最大和最小值实例
2020/04/02 Python
Pytorch如何切换 cpu和gpu的使用详解
2021/03/01 Python
机械绘图员岗位职责
2013/11/19 职场文书
十佳教师事迹材料
2014/01/11 职场文书
学年末自我鉴定
2014/01/21 职场文书
C++程序员求职信范文
2014/04/14 职场文书
快消品行业营销模式与盈利模式分享
2019/09/27 职场文书
nginx中proxy_pass各种用法详解
2021/11/07 Servers