pytorch将pt模型转onnx模型

pytorch将pt模型转onnx模型

一 导出ONNX模型

torch.onnx.export(
model,  # 要导出的模型
args, #  模型的输入参数,输入参数只需满足shape正确
onnx_export_filepath, # 转换输出的onnx模型的路径
export_params=True, # true表示导出trained model，否则untrained model。默认即可
verbose=False, #  true表示打印调试信息
input_names=None, # 指定输入节点名称
output_names=None, # 指定输出节点名称
do_constant_folding=True, # 是否使用常量折叠，默认即可
dynamic_axes=None, # 通过dynamic_axes来指定输入tensor的哪些参数可变
opset_version=9 # 指定onnx的opset版本，版本过低的话，不支持upsample等操作
)

示例1：

import torch
model = torch.load('HyClsNet(traced).pt')
model.eval()
input_names = ['input']
output_names = ['output']
x = torch.randn(1,3,512,512,requires_grad=True)
torch.onnx.export(model, x, 'HyClsNet(traced).onnx', input_names=input_names, output_names=output_names, verbose='True')

示例2：gpu版模型

import torch
import torch.nn
import onnx
import numpy as np
torch.backends.cudnn.benchmark = False


model = torch.load('HyClsNet(traced).pt')
model.eval()
model.cuda()
input_names = ['input']
output_names = ['output']
x = torch.randn(1,3,512,512).cuda()
torch.onnx.export(model, x, 'HyClsNet(traced).onnx', input_names=input_names, output_names=output_names, verbose='True')

二 Pytorch使用onnx模型安装包，纯导出，不需要

CPU

pip install onnx -i http://pypi.douban.com/simple/  --trusted-host pypi.douban.com
pip install onnxruntime -i http://pypi.douban.com/simple/  --trusted-host pypi.douban.com

GPU

pip install onnx -i http://pypi.douban.com/simple/  --trusted-host pypi.douban.com
pip install onnxruntime-gpu  -i http://pypi.douban.com/simple/  --trusted-host pypi.douban.com

如果安装了CPU版的onnxruntime，要先pip uninstall onnxruntime，再安装对应的GPU版本

三 使用ONNX模型

import onnx
import onnxruntime as ort
model = onnx.load('HyClsNet(traced).onnx')
onnx.checker.check_model(model)
session = ort.InferenceSession('HyClsNet(traced).onnx')
x=np.random.randn(1,3,512,512).astype(np.float32)  # 输入type是np.float32
outputs = session.run(None,input = {'input' : x })