PyTorch:模型推理加速之onnx
Open Neural Network Exchange(ONNX,开放神经网络交换)格式,是一个用于表示深度学习模型的标准,可使模型在不同框架之间进行转移。ONNX是一种针对机器学习所设计的开放式的文件格式,用于存储训练好的模型。它使得不同的人工智能框架(如Pytorch, MXNet)可以采用相同格式存储模型数据并交互。ONNX的规范及代码主要由微软,亚马逊 ,Facebook 和 IBM 等
简介
模型部署流水线
为了让模型最终能够部署到某一环境上,开发者们可以使用任意一种深度学习框架来定义网络结构,并通过训练确定网络中的参数。之后,模型的结构和参数会被转换成一种只描述网络结构的中间表示,一些针对网络结构的优化会在中间表示上进行。最后,用面向硬件的高性能编程框架(如 CUDA,OpenCL)编写,能高效执行深度学习网络中算子的推理引擎会把中间表示转换成特定的文件格式,并在对应硬件平台上高效运行模型。
Open Neural Network Exchange(ONNX,开放神经网络交换)格式:是一个用于表示深度学习模型的标准,可使模型在不同框架之间进行转移。ONNX是一种针对机器学习所设计的开放式的文件格式,用于存储训练好的模型。它使得不同的人工智能框架(如Pytorch, MXNet)可以采用相同格式存储模型数据并交互。 ONNX的规范及代码主要由微软,亚马逊 ,Facebook 和 IBM 等公司共同开发,以开放源代码的方式托管在Github上。目前官方支持加载ONNX模型并进行推理的深度学习框架有: Caffe2, PyTorch, MXNet,ML.NET,TensorRT 和 Microsoft CNTK,并且 TensorFlow 也非官方的支持ONNX。---维基百科
ONNX是一个开放式规范,由以下组件组成:可扩展计算图模型的定义;标准数据类型的定义;内置运算符的定义
推理引擎 -ONNX Runtime:ONNX Runtime 是由微软维护的一个跨平台机器学习推理加速器,也就是我们前面提到的”推理引擎“。ONNX Runtime 是直接对接 ONNX 的,即 ONNX Runtime 可以直接读取并运行 .onnx 文件, 而不需要再把 .onnx 格式的文件转换成其他格式的文件。
安装依赖
$pip install optimum[onnxruntime]
Note: 也可以通过下面安装必须的
pip install optimum
pip install onnx
pip install onnxruntime # CPU build
pip install onnxruntime-gpu # GPU build
上面会安装optimum onnx onnxruntime,并多出 evaluate responses。
Onnx模型导出
torch.onnx.export(model, args, f, export_params=True, verbose=False, training=<TrainingMode.EVAL: 0>, input_names=None, output_names=None, operator_export_type=<OperatorExportTypes.ONNX: 0>, opset_version=None, do_constant_folding=True, dynamic_axes=None, keep_initializers_as_inputs=None, custom_opsets=None, export_modules_as_functions=False)
参数
- model:要转换的torch模型
- args:输入,即dummy_input,就是一个输入的实例,仅提供输入shape、type等信息,里面什么内容不重要,所以用randn, zeros, ones都是可以的,使用真实输入的一个case转换后作为输入最好; 作用应该是为了让onnx探查模型forward时input和output的关系及经过的必要的图吧。args格式一般是 (tuple or torch.Tensor)但是试了dict也可以。另外一般pytorch模型的输入为tensor。
- f:保存的onnx文件名。
- export_params=True:模型中是否存储模型权重。一般中间表示包含两大类信息:模型结构和模型权重,这两类信息可以在同一个文件里存储,也可以分文件存储。ONNX 是用同一个文件表示记录模型的结构和权重的。
我们部署时一般都默认这个参数为 True。如果 onnx 文件是用来在不同框架间传递模型(比如 PyTorch 到 Tensorflow)而不是用于部署,则可以令这个参数为 False。 - verbose=False:verbose=True causes the exporter to print out a human-readable representation of the model.if True, prints a description of the model being exported to stdout. In addition, the final ONNX graph will include the field doc_string` from the exported model which mentions the source code locations for model. If True, ONNX exporter logging will be turned on. 但是这样模型大小会变大,比如原本20m变成30m。
- input_names=None, output_names=None:– names to assign to the input/output nodes of the graph, in order. input_names没有特别要求,可以自己定义,但是要与session里输入的一致。其中output_names在session.run时候可以不写,直接写成None。
- opset_version=None:opset_version (int, default 14) – The version of the default (ai.onnx) opset to target. Must be >= 7 and <= 16.指定的操作版本,一般越高的版本会支持更多的操作。如果遇到某个操作不支持,可以将版本号设置的高一点试试。
- do_constant_folding=True:是否执行常量折叠优化。Constant-folding will replace some of the ops that have all constant inputs with pre-computed constant nodes.
- dynamic_axes=None:dynamic_axes={"input":{0:"batch_size"}, "output":{0:"batch_size"}})。
dummy_input中不要输入None,它会将后面数据的size搞过来,导致后面的size错位而出错以及run推理时候出错。可以改成空的一个可替代值,如oov_zeros = torch.zeros((batch_size, 0), dtype=torch.float32)。另外export时,forward时输入dummy_input中的常量参数会被转成tensor,如beam_size: tensor(5),但是None还是保持不变就是None(可能会出问题)。
input_names变量名一定要和dummy_input保持一致(至少是前面顺序的),写错位置就相当于将输入的dummy_input交换位置输入到模型中了。(后面的)input_names里面没写的,dummy_input有的且没有被常量折叠的,在export时,也是会正常过model.forward,而且会给你在input_names中加一个名字,如'inp'。
opset_version可以不设置。如果设置小了,比如opset_version=11还会出问题,比如高级操作einsum无法转换。
动态维度
如果不设置动态维度,那么dummy_input是什么维度,真实输入也必须是,但是这样变长的输入就不可行了,一般解决办法就是加入动态维度。动态维度dynamic_axes两种设置方式:一般用列表的方式表示动态维度,例如:dynamic_axes_0 = { 'in' : [0], 'out' : [0] } 。由于 ONNX 要求每个动态维度都有一个名字,这样写的话会引出一条UserWarning,警告我们通过列表的方式设置动态维度的话系统会自动为它们分配名字。一种显式添加动态维度名字的方法如下:dynamic_axes_0 = {'in' : {0: 'batch'}, 'out' : {0: 'batch'} }
多输入
比如model def forward(self, x, scale):多出一个参数scale,我们可以通过将参数列表包在一个list中,如args=(x, 2)将dummy_input变成一个list类的。或者也可以像示例3中使用dict。
Note: 由于export函数的机制,会把模型输入的参数自动转换成tensor类型,比如上面的scale参数,虽然传入的时候是int32类型,但是export在执行时会调用到forward函数,此时scale已经变成一个tensor类型。[支持多参数与动态输入]
输出
原始输出如果是一个tensor的list,那么实际onnx会认为是有很多个tensor输出。这时你不指定output_names时,onnx自动生成的output_names = ['inp', 'inp.3',...],且定长(dummy_input对应长度)。
如果你只指定output_names=['output'],onnx会自动生成['output', 'inp', ...]这种。当你outputs = ort_session.run(['output'], ort_inputs)只指定了一个'output'时,相当于只会运行第一个输出,这时返回的结果就只有一个字符。当然你可以这样:outputs = ort_session.run([output.name for output in ort_session._outputs_meta], ort_inputs)或者outputs = ort_session.run(None, ort_inputs),但是这个还是定长的。
Note: 如果模型推理部分存在while循环,要全部转成onnx是不行的,因为onnx完全是静态图。比如你导出时的while循环了k次,后面不管什么query都是k次,即使使用了动态维度,因为他的onnx图可以看到,最后直接生成的是死的k个输出(或者再汇总到最后一个输出),但是这k个是固定死了。
The resulting .onnx file contains a binary protocol buffer which contains both the network structure and parameters of the model you exported.
示例
示例1:
class SumModule(torch.nn.Module):
def forward(self, x):
return torch.sum(x, dim=1)
torch.onnx.export(
SumModule(),
(torch.ones(2, 2),),
"onnx.pb",
input_names=["x"],
output_names=["sum"]
)
示例2:
input_name = 'input'
output_name = 'output'
torch.onnx.export(model,
input_data,
"simplenet.onnx",
opset_version=11,
input_names=[input_name],
output_names=[output_name],
dynamic_axes={
input_name: {0: 'batch_size', 2: 'in_width', 3: 'int_height'},
output_name: {0: 'batch_size', 2: 'out_width', 3: 'out_height'}}
)
示例3:
with torch.no_grad():
model = ***
model.load_state_dict(***)
model.eval()
dummy_content = '这是一个句子'
dummy_encoder_input = torch.tensor([[vocab.word_2_idx(_) for _ in dummy_content]])
dummy_input = {'encoder_input': dummy_encoder_input,
'encoder_mask': torch.ones_like(dummy_encoder_input),
'context_vec': torch.randn(1, hidden_dim * 2),
'coverage': torch.randn(1, len(dummy_content)),
'oovs_zero': torch.zeros((1, 0), dtype=torch.float32),
'beam_size': beam_size, 'mode': "predict", 'vocab': None}
torch.onnx.export(model, dummy_input, onnx_filename,
input_names=['encoder_input', 'encoder_mask', 'encoder_with_oov',
'context_vec', 'coverage', 'oovs_zero', 'beam_size'],
output_names=['output'],
dynamic_axes={
'encoder_input': {1: 'seq_len'},
'encoder_mask': {1: 'seq_len'},
'coverage': {1: 'seq_len'},
'oovs_zero': {1: 'oovs_len'},
'output': {0: 'out_seq_len'}
}
)
bug fix
1 RuntimeError: Only tuples, lists and Variables are supported as JIT inputs/outputs. Dictionaries and strings are also accepted, but their usage is not recommended. Here, received an input of unsupported type: Beam
原因:模型forward()输入输出只能是这些类型。如果说,比如forward()输出返回一个自定义的Beam对象,那就会报上面错误。
解决:输出Beam对象改成输出对象的某个需要输出的属性,比如.tokens就是返回它的一个list子属性。
2 RuntimeError: output 1 (644[ CPULongType{} ]) of traced region did not have observable data dependence with trace inputs; this probably indicates your program cannot be understood by the tracer.
原因:输出return sorted_beams_result[0].tokens[1:],如果这个tokens中的元素加入时,过了.item()转换成了数字,没法trace。官网也有写到
Avoid NumPy and built-in Python types (PyTorch models can be written using NumPy or Python types and functions, but during tracing, any variables of NumPy or Python types (rather than torch.Tensor) are converted to constants, which will produce the wrong result if those values should change depending on the inputs.
And rather than use torch.Tensor.item() (which converts a Tensor to a Python built-in number):
# Bad! y.item() will be replaced with a constant during tracing.
def forward(self, x, y):
return x.reshape(y.item(), -1)
);
Avoid Tensor.data;
Avoid in-place operations when using tensor.shape in tracing mode]
有时候onnx模型的input_names = []也是这个原因,即输出前一步都是.item后的数据,然后再变成tensor输出的话,input是不能trace到output的。
解决:所以去掉转换,保持原有的tensor格式就ok。
3 RuntimeError: "Dynamic shape axis should be no more than the shape dimension for" seq_len
原因:如果你以为原始输出是一个2维tensor(实际是多个1维tensor的lsit),而将output维度1设置成动态的就可能出这个错误。
4 TracerWarning: torch.tensor results are registered as constants in the trace.
貌似不能直接用torch.tensor()来转换tensor,会导致当成常量而无法trace。
5 TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
要用torch的比较,但是torch比较需要tensor,但又不能直接用torch.tensor()来转换tensor,会导致当成常量而无法trace。所以一般可以不管?
Onnx模型加载和推理
验证onnx模型正确性
verify the output
import onnx
# Load the ONNX model
model = onnx.load("alexnet.onnx")
# Check that the model is well formed. Check the consistency of a model. An exception is raised if the test fails
onnx.checker.check_model(model)
# Print a human readable representation of the graph 加了verbose=True会多输出一些东西,但是不多。
print(onnx.helper.printable_graph(model.graph))
查看onnx模型输入输出是什么样的数据
import onnx
from onnx import mapping
model = onnx.load_model('model.onnx')
# 获取输入节点
input_nodes = model.graph.input
# 获取输出节点
output_nodes = model.graph.output
# 遍历输入节点并查看属性
for node in input_nodes:
print(f"Input name: {node.name}")
print(f"Input shape: {node.type.tensor_type.shape.dim}")
print(f"Input data type: {node.type.tensor_type.elem_type}")
print(f"Input data type: {mapping.TENSOR_TYPE_TO_NP_TYPE[node.type.tensor_type.elem_type]}")
其中node.type.tensor_type.elem_type返回的是一个数字,而不是直接的数据类型。这是因为ONNX规范中使用了预定义的整数值来表示不同的数据类型:1:FLOAT;2:UINT8;3:INT8;4:UINT16;5:INT16;6:INT32;7:INT64;8:STRING;9:BOOL;10:FLOAT16;11:DOUBLE;12:UINT32;13:UINT64;14:COMPLEX64;15:COMPLEX128
结果大概是这样的:分别是动态tensor、静态tensor、int的dummy_input,以及动态tensor输出,所导出模型的输入输出情况。
Input name: encoder_input
Input shape: [dim_value: 1, dim_param: "seq_len"]
Input data type: int64
Input name: context_vec
Input shape: [dim_value: 1, dim_value: 512]
Input data type: float32
Input name: beam_size
Input shape: []
Input data type: int64
Input name: output
Input shape: [dim_param: "out_seq_len"]
Input data type: int64
Note: 动态维度不指定名字时,自动取名是这样:encoder_input_dynamic_axes_1、output_dynamic_axes_1。
也可以在infersession加载后查看所有输入输出名:
ort_session = ort.InferenceSession(onnx_filename)
print('input_names = ', [i.name for i in ort_session.get_inputs()])
print('output_names = ', [output.name for output in ort_session._outputs_meta])
输出:input_names = ['encoder_input', ...'coverage', 'oovs_zero', 'beam_size']
output_names = ['output']
使用Netron
输出的模型使用Netron打开,
查看输入输出信息可以看到动态维度不是数字了。如输入的维度变成:[batch_size,1,in_width,int_height],输出的维度变成:[batch_size,1,out_width,out_height]。表示这个模型可以接收动态的批次大小和宽高尺寸。
load and run the model
[(OPTIONAL) EXPORTING A MODEL FROM PYTORCH TO ONNX AND RUNNING IT USING ONNX RUNTIME]
示例1:
import onnxruntime as ort
ort_session = ort.InferenceSession("alexnet.onnx")
outputs = ort_session.run(
None,
{"actual_input_1": np.random.randn(10, 3, 224, 224).astype(np.float32)},
)
print(outputs[0])
示例2:
ort_session = ort.InferenceSession(onnx_filename)
def to_numpy(tensor):
return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
# compute ONNX Runtime output prediction
import numpy as np
ort_inputs = {'encoder_input': to_numpy(batch[0]),
'encoder_mask': to_numpy(batch[1]),
'encoder_with_oov': to_numpy(batch[2]),
'context_vec': to_numpy(batch[4]),
'coverage': to_numpy(batch[5]),
'oovs_zero': to_numpy(batch[3]),
'beam_size': np.asarray([model_config.beam_size], dtype=np.int64),
}
# outputs = ort_session.run([output.name for output in ort_session._outputs_meta], ort_inputs)
# outputs = ort_session.run(['output'], ort_inputs)
outputs = ort_session.run(None, ort_inputs)
hypothesis_idx_list = outputs[0]
输入
onnx的输入只能是array类型。
输出
run时候,如果output_names设置成None,就是默认所有的都输出。指定的话,如果少指定了,就只返回指定的。
返回一个list,即将所有返回的output(tensor)转成array,添加到list中返回。
所以如果所有输出放到一个tensor中,取outputs[0]即可;但是放到n个tensor中,那list中有很多个array。
bug fix
1 INVALID_ARGUMENT : Unexpected input data type. Actual: (tensor(float)) , expected: (tensor(int64))?
一般都是dummy_input和加载后推理时真实输入不同导致,要么是前面说到的输入位置错了,要么就是类型没对上。
2 ort_session.run(['output'], ort_inputs)出错:onnxruntime\core/framework/op_kernel_context.h:42 onnxruntime::OpKernelContext::Input Missing Input: oovs_zero
原因:应该是输入oovs_zero对应的值是None的话,就出错。
解决:改成空的一个可替代值,如oov_zeros = torch.zeros((batch_size, 0), dtype=torch.float32)
3 ort_session.run(['output'], ort_inputs)出错:ValueError: Required inputs (['mode']) are missing from input feed ([...]).
原因:有时,模型在export为onnx时,即使input_names和dummy_input里面都加了某个参数如mode='predict',模型export后可能会删除(应该是当成常量优化了),但是有时可能又和你加的其它输入相关这样就又不会被删除。
解决:这时需要查看一下onnx的输入,然后把真实的input里面加入之前本来不需要加的。
4 ort_session.run(['output'], ort_inputs)出错:RuntimeError: Unable to handle object of type <class 'int'>
原因:onnx真实输入ort_inputs只接受ndarray类型值,int和tensor都不行。
解决:{'beam_size': np.asarray([model_config.beam_size], dtype=np.int64)}
5 onnxruntime.capi.onnxruntime_pybind11_state.InvalidArgument: [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Got invalid dimensions for input: coverage for the following indices index: 1 Got: 34 Expected: 25
原因:export时输入dummy_input维度1(一般是seq_len)是定长的25,真实输入时又是34长度。
解决:将输入的维度1加到dynamic_axes中指定为动态可变长。
6 ort_session.run时出错:INVALID_ARGUMENT : Invalid Feed Input Name:context_vec。
原因:检查发现onnx模型的input_names = []。原因是export时,用了.item未trace到。
解决:所以去掉转换,保持原有的tensor格式就ok。
from: -柚子皮-
ref: [TORCH.ONNX]***
[Pytorch分类模型转onnx以及onnx模型推理 - 知乎]*
未尝试的[onnxsim]
开放原子开发者工作坊旨在鼓励更多人参与开源活动,与志同道合的开发者们相互交流开发经验、分享开发心得、获取前沿技术趋势。工作坊有多种形式的开发者活动,如meetup、训练营等,主打技术交流,干货满满,真诚地邀请各位开发者共同参与!
更多推荐
所有评论(0)