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环境下搭建属于自己的pip源的教程
May 05 Python
django之session与分页(实例讲解)
Nov 13 Python
python实现换位加密算法的示例
Oct 14 Python
pycharm 将python文件打包为exe格式的方法
Jan 16 Python
如何在Django中设置定时任务的方法示例
Jan 18 Python
python用requests实现http请求代码实例
Oct 31 Python
pytorch中的上采样以及各种反操作,求逆操作详解
Jan 03 Python
使用Python将Exception异常错误堆栈信息写入日志文件
Apr 08 Python
Jupyter Notebook远程登录及密码设置操作
Apr 10 Python
关于jupyter打开之后不能直接跳转到浏览器的解决方式
Apr 13 Python
Python实现日志实时监测的示例详解
Apr 06 Python
PyTorch中permute的使用方法
Apr 26 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
深入浅出php socket编程
2015/05/13 PHP
利用ajax和PHP实现简单的流程管理
2017/03/23 PHP
浅谈Yii乐观锁的使用及原理
2017/07/25 PHP
PHP7中I/O模型内核剖析详解
2019/04/14 PHP
JS location几个方法小姐
2008/07/09 Javascript
Javascript 篱式条件判断
2008/08/22 Javascript
jQuery 菜单随滚条改为以定位方式(固定要浏览器顶部)
2012/05/24 Javascript
通过js来制作复选框的全选和不选效果
2014/05/22 Javascript
JavaScript中实现依赖注入的思路分享
2015/01/15 Javascript
Jquery 实现grid绑定模板
2015/01/28 Javascript
JavaScript数据结构与算法之栈与队列
2016/01/29 Javascript
AngularJS基础 ng-show 指令简单示例
2016/08/03 Javascript
jQuery Validate验证表单时多个name相同的元素只验证第一个的解决方法
2016/12/24 Javascript
微信小程序  http请求封装详解及实例代码
2017/02/15 Javascript
详解Angular 4.x 动态创建组件
2017/04/25 Javascript
js中的事件委托或是事件代理使用详解
2017/06/23 Javascript
微信小程序三级联动地址选择器的实例代码
2017/07/12 Javascript
基于Vue 实现一个中规中矩loading组件
2019/04/03 Javascript
vue输入框使用模糊搜索功能的实现代码
2020/05/26 Javascript
[02:19]DOTA选手解说齐贺岁
2018/02/11 DOTA
[01:36]极致酷炫!TI9典藏宝瓶+撼地者至宝展示
2019/06/11 DOTA
python调用百度语音识别实现大音频文件语音识别功能
2018/08/30 Python
Python英文文本分词(无空格)模块wordninja的使用实例
2019/02/20 Python
python-视频分帧&amp;多帧合成视频实例
2019/12/10 Python
python3.8.3安装教程及环境配置的详细教程(64-bit)
2020/11/28 Python
python 基于Apscheduler实现定时任务
2020/12/15 Python
HTML5 Canvas阴影使用方法实例演示
2013/08/02 HTML / CSS
英国乐购杂货:Tesco Groceries
2018/11/29 全球购物
高二化学教学反思
2014/01/30 职场文书
cf收人广告词大全
2014/03/14 职场文书
励志演讲稿300字
2014/08/21 职场文书
2014年小学重阳节活动策划方案
2014/09/16 职场文书
汇报材料怎么写
2014/12/30 职场文书
文艺演出主持词
2015/07/01 职场文书
HDFS免重启挂载新磁盘
2022/04/06 Servers
Python中re模块的元字符使用小结
2022/04/07 Python