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