官方资料
https://pytorch.org/docs/stable/onnx.html#id2
onnx ops
https://github.com/onnx/onnx/blob/master/docs/Operators.md
pytorch onnx
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = torch.load("test.pth") # pytorch模型加载
batch_size = 1 #批处理大小
input_shape = (3, 244, 384) #输入数据,改成自己的输入shape
# #set the model to inference mode
model.eval()
x = torch.randn(batch_size, *input_shape) # 生成张量
x = x.to(device)
export_onnx_file = "test.onnx" # 目的ONNX文件名
torch.onnx.export(model
x,
export_onnx_file,
opset_version=10,
do_constant_folding=True, # 是否执行常量折叠优化
input_names=["input"], # 输入名
output_names=["output"], # 输出名
dynamic_axes={"input":{0:"batch_size"}, # 批处理变量
"output":{0:"batch_size"}})
接口参数
torch.onnx.export(model, args, f,
export_params=True,
verbose=False,
training=<TrainingMode.EVAL: 0>,
input_names=None,
output_names=None,
aten=False,
export_raw_ir=False,
operator_export_type=None,
opset_version=None,
_retain_param_name=True,
do_constant_folding=True,
example_outputs=None,
strip_doc_string=True,
dynamic_axes=None,
keep_initializers_as_inputs=None,
custom_opsets=None,
enable_onnx_checker=True,
use_external_data_format=False)
[SOURCE]
- model (torch.nn.Module) – the model to be exported. 输入模型
- args (tuple of arguments or torch.Tensor, a dictionary consisting of named arguments (optional)) –
a dictionary to specify the input to the corresponding named parameter: - KEY: str, named parameter - VALUE: corresponding input args can be structured either as:
输入数据参数
- f – a file-like object (has to implement fileno that returns a file descriptor) or a string containing a file name. A binary Protobuf will be written to this file. 输出文件名
- export_params (bool, default True) – if specified, all parameters will be exported. Set this to False if you want to export an untrained model. In this case, the exported model will first take all of its parameters as arguments, the ordering as specified by
model.state_dict().values() True 导出参数
- verbose (bool, default False) – if specified, we will print out a debug description of the trace being exported. 输出debug信息
- training (enum, default TrainingMode.EVAL) – TrainingMode.EVAL: export the model in inference mode. TrainingMode.PRESERVE: export the model in inference mode if model.training is False and to a training friendly mode if model.training is True. TrainingMode.TRAINING: export the model in a training friendly mode.
- input_names (list of strings, default empty list) – names to assign to the input nodes of the graph, in order
输入参数名
- output_names (list of strings, default empty list) – names to assign to the output nodes of the graph, in order
- aten (bool, default False) – [DEPRECATED. use operator_export_type] export the model in aten mode. If using aten mode, all the ops original exported by the functions in symbolic_opset
.py are exported as ATen ops. - export_raw_ir (bool, default False) – [DEPRECATED. use operator_export_type] export the internal IR directly instead of converting it to ONNX ops.
- operator_export_type (enum, default OperatorExportTypes.ONNX) –
OperatorExportTypes.ONNX: All ops are exported as regular ONNX ops (with ONNX namespace). OperatorExportTypes.ONNX_ATEN: All ops are exported as ATen ops (with aten namespace). OperatorExportTypes.ONNX_ATEN_FALLBACK: If an ATen op is not supported in ONNX or its symbolic is missing, fall back on ATen op.
ONNX ops 不支持 pytorch/ATen ops 时如何操作。
https://www.onnxruntime.ai/docs/how-to/export-pytorch-model.html
- opset_version (int, default is 9) – by default we export the model to the opset version of the onnx submodule. Since ONNX’s latest opset may evolve before next stable release, by default we export to one stable opset version. Right now, supported stable opset version is 9. The opset_version must be _onnx_main_opset or in _onnx_stable_opsets which are defined in torch/onnx/symbolic_helper.py
我们默认导出到一个稳定的opset版本。现在,支持稳定的opset版本是9。
https://github.com/onnx/onnx/blob/master/docs/Versioning.md onnx 版本
https://github.com/microsoft/onnxruntime/blob/master/docs/Versioning.md onnxruntime/onnx 版本
- do_constant_folding (bool, default False) – If True, the constant-folding optimization is applied to the model during export. Constant-folding optimization will replace some of the ops that have all constant inputs, with pre-computed constant nodes.如果为真,则在导出过程中对模型应用常数折叠优化。常数折叠优化将用预先计算的常数节点替换一些具有所有常数输入的ops。
- example_outputs (tuple of Tensors, default None) – Model’s example outputs being exported. example_outputs must be provided when exporting a ScriptModule or TorchScript Function.
被导出的模型示例输出。当导出ScriptModule或TorchScript函数时,必须提供example_outputs。
- strip_doc_string (bool, default True) – if True, strips the field “doc_string” from the exported model, which information about the stack trace.
- dynamic_axes (dict
> or dict<string, list(int)>, default empty dict) –
a dictionary to specify dynamic axes of input/output, such that: - KEY: input and/or output names - VALUE: index of dynamic axes for given key and potentially the name to be used for exported dynamic axes. In general the value is defined according to one of the following ways or a combination of both: (1). A list of integers specifying the dynamic axes of provided input. In this scenario automated names will be generated and applied to dynamic axes of provided input/output during export. OR (2). An inner dictionary that specifies a mapping FROM the index of dynamic axis in corresponding input/output TO the name that is desired to be applied on such axis of such input/output during export.
指定输入/输出的动态轴的字典
- keep_initializers_as_inputs (bool, default None) –
If True, all the initializers (typically corresponding to parameters) in the exported graph will also be added as inputs to the graph. If False, then initializers are not added as inputs to the graph, and only the non-parameter inputs are added as inputs.
This may allow for better optimizations (such as constant folding etc.) by backends/runtimes that execute these graphs. If unspecified (default None), then the behavior is chosen automatically as follows. If operator_export_type is OperatorExportTypes.ONNX, the behavior is equivalent to setting this argument to False. For other values of operator_export_type, the behavior is equivalent to setting this argument to True. Note that for ONNX opset version < 9, initializers MUST be part of graph inputs. Therefore, if opset_version argument is set to a 8 or lower, this argument will be ignored.
如果为True,则导出图形中的所有初始值设定项(通常对应于参数)也将作为输入添加到图形中。如果为False,则不将初始值设定项作为输入添加到图中,而仅将非参数输入作为输入添加。
这可能允许通过执行这些图的后端/运行时进行更好的优化(例如常量折叠等)。如果未指定(默认无),则自动选择行为,如下所示。如果运算符\导出\类型为运算符导出类型.ONNX,该行为相当于将此参数设置为False。对于运算符\导出\类型的其他值,该行为等效于将此参数设置为True。注意,对于ONNX opset version<9,初始化器必须是图形输入的一部分。因此,如果opset\u version参数设置为8或更低,则将忽略此参数。
- custom_opsets (dict
, default empty dict ) – A dictionary to indicate custom opset domain and version at export. If model contains a custom opset, it is optional to specify the domain and opset version in the dictionary: - KEY: opset domain name - VALUE: opset version If the custom opset is not provided in this dictionary, opset version is set to 1 by default. - 在导出时指示自定义opset域和版本的字典。—KEY: opset域名—VALUE: opset版本如果本字典中没有自定义的opset,则opset版本默认为1。
- enable_onnx_checker (bool, default True) – If True the onnx model checker will be run as part of the export, to ensure the exported model is a valid ONNX model.
- external_data_format (bool, default False) – If True, then the model is exported in ONNX external data format, in which case some of the model parameters are stored in external binary files and not in the ONNX model file itself.
3.onnxruntime
```
import os, sys sys.path.append(os.getcwd()) import onnxruntime import onnx import cv2 import torch import numpy as np import torchvision.transforms as transforms
class ONNXModel(): def init(self, onnx_path): “”” :param onnx_path: “”” self.onnx_session = onnxruntime.InferenceSession(onnx_path) self.input_name = self.get_input_name(self.onnx_session) self.output_name = self.get_output_name(self.onnx_session) print(“input_name:{}”.format(self.input_name)) print(“output_name:{}”.format(self.output_name))
def get_output_name(self, onnx_session):
"""
output_name = onnx_session.get_outputs()[0].name
:param onnx_session:
:return:
"""
output_name = []
for node in onnx_session.get_outputs():
output_name.append(node.name)
return output_name
def get_input_name(self, onnx_session):
"""
input_name = onnx_session.get_inputs()[0].name
:param onnx_session:
:return:
"""
input_name = []
for node in onnx_session.get_inputs():
input_name.append(node.name)
return input_name
def get_input_feed(self, input_name, image_numpy):
"""
input_feed={self.input_name: image_numpy}
:param input_name:
:param image_numpy:
:return:
"""
input_feed = {}
for name in input_name:
input_feed[name] = image_numpy
return input_feed
def forward(self, image_numpy):
'''
# image_numpy = image.transpose(2, 0, 1)
# image_numpy = image_numpy[np.newaxis, :]
# onnx_session.run([output_name], {input_name: x})
# :param image_numpy:
# :return:
'''
# 输入数据的类型必须与模型一致,以下三种写法都是可以的
# scores, boxes = self.onnx_session.run(None, {self.input_name: image_numpy})
# scores, boxes = self.onnx_session.run(self.output_name, input_feed={self.input_name: iimage_numpy})
input_feed = self.get_input_feed(self.input_name, image_numpy)
scores, boxes = self.onnx_session.run(self.output_name, input_feed=input_feed)
return scores, boxes
def to_numpy(tensor): return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
r_model_path=”/home/zigangzhao/DMS/mtcnn-pytorch/test0815/onnx_model/rnet.onnx” o_model_path=”/home/zigangzhao/DMS/mtcnn-pytorch/test0815/onnx_model/onet.onnx”
img = cv2.imread(“/home/zigangzhao/DMS/mtcnn-pytorch/data_set/train/24/positive/999.jpg”) img = cv2.resize(img, 24, 24), interpolation=cv2.INTER_CUBIC)
totensor = transforms.ToTensor() img = to_tensor(img) img = img.unsqueeze(0)
方法1:
rnet1 = ONNXModel(r_model_path) out = rnet1.forward(to_numpy(img))
print(out)
方法2: rnet_session = onnxruntime.InferenceSession(r_model_path) onet_session = onnxruntime.InferenceSession(o_model_path)
compute ONNX Runtime output prediction
inputs = {onet_session.get_inputs()[0].name: to_numpy(img)} outs = onet_session.run(None, inputs) img_out_y = outs
print(img_out_y)
```