将torch模型转化为onnx模型

函数

torch.onnx.export(
model, 
args, 
f, 
export_params=True, 
verbose=False, 
training=False, 
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=False,
example_outputs=None, 
strip_doc_string=True, 
dynamic_axes=None, 
keep_initializers_as_inputs=None)

功能:

将模型以ONNX格式导出并保存.

参数

• model (torch.nn.Module) – 要导出的模型.
• args (tuple of arguments) – 模型的输入, 任何非Tensor参数都将硬编码到导出的模型中；任何Tensor参数都将成为导出的模型的输入，并按照他们在args中出现的顺序输入。因为export运行模型，所以我们需要提供一个输入张量x。只要是正确的类型和大小，其中的值就可以是随机的。请注意，除非指定为动态轴，否则输入尺寸将在导出的ONNX图形中固定为所有输入尺寸。在此示例中，我们使用输入batch_size 1导出模型，但随后dynamic_axes 在torch.onnx.export()。因此，导出的模型将接受大小为[batch_size，3、100、100]的输入，其中batch_size可以是可变的。
• export_params (bool, default True) – 如果指定为True或默认, 参数也会被导出. 如果你要导出一个没训练过的就设为 False
• verbose (bool, default False) - 如果指定，我们将打印出一个导出轨迹的调试描述。
• training (bool, default False) - 在训练模式下导出模型。目前，ONNX导出的模型只是为了做推断，所以你通常不需要将其设置为True
• input_names (list of strings, default empty list) – 按顺序分配名称到图中的输入节点
• output_names (list of strings, default empty list) –按顺序分配名称到图中的输出节点
• dynamic_axes – {‘input’ : {0 : ‘batch_size’}, ‘output’ : {0 : ‘batch_size’}}) # variable lenght axes
  举例:

shape(input_1) = ('b', 3, 'w', 'h')
and shape(input_2) = ('b', 4)
and shape(output)  = ('b', 'd', 5)
torch.onnx.export(model, args, f, export_params=True, verbose=False, training=False, input_names=None, output_names=None, dynamic_axes = {‘input_1’:[0, 2, 3], ‘input_2’:[0], ‘output’:[0, 1]})

首先加载我们的模型，例如加载我的模型：

model = ResNetSE34V2(8222, device, encoder_type='ASP').to(device)
model = load_checkpoint(model, '../checkpoint/', checkpoint_name='model1.pt')
model.eval()

加载完成后一定要让模型进入eval状态 ，之后就可以进行导出操作了：

x = torch.randn((1, 80, 300))
torch.onnx.export(model,  # model being run
				  x,  # model input (or a tuple for multiple inputs)
				  "model.onnx",  # where to save the model (can be a file or file-like object)
				  export_params=True,  # store the trained parameter weights inside the model file
				  opset_version=10,  # the ONNX version to export the model to
				  do_constant_folding=True,  # whether to execute constant folding for optimization
				  input_names=['input'],  # the model's input names
				  output_names=['output']  # the model's output names)

动态shape问题

如此即可导出模型，但如果我们的输入模型参数是不断变化的呢，例如第三个维度的长度不是300，又该如何导出？我们只需要指定一下 第二个维度为动态变化的即可 ，如下

torch.onnx.export(model,  # model being run
				  x,  # model input (or a tuple for multiple inputs)
				  "model2.onnx",  # where to save the model (can be a file or file-like object)
				  export_params=True,  # store the trained parameter weights inside the model file
				  opset_version=10,  # the ONNX version to export the model to
				  do_constant_folding=True,  # whether to execute constant folding for optimization
				  input_names=['input'],  # the model's input names
				  output_names=['output'],  # the model's output names
				  dynamic_axes={'input': {2: 'int_height'}})

那如果有多个变化的，也是写上相应维度即可，并且输出也可以是变化的，如下：

torch.onnx.export(model,  # model being run
				  x,  # model input (or a tuple for multiple inputs)
				  "model2.onnx",  # where to save the model (can be a file or file-like object)
				  export_params=True,  # store the trained parameter weights inside the model file
				  opset_version=10,  # the ONNX version to export the model to
				  do_constant_folding=True,  # whether to execute constant folding for optimization
				  input_names=['input'],  # the model's input names
				  output_names=['output'],  # the model's output names
				  dynamic_axes= {
                       input_name: {0: 'batch_size', 2 : 'in_width', 3: 'int_height'},
                       output_name: {0: 'batch_size', 2: 'out_width', 3:'out_height'}})

模型导出完成，实际使用一般使用cpu进行infer，此处下载cpu版本的onnx，不能同时下载cpu和gpu版本！

pip install onnxruntime

使用模型：

x = torch.randn((1, 80, 300))
onnx_model = onnxruntime.InferenceSession("model2.onnx")
print(onnx_model.get_inputs()[0].name)
inputs = {onnx_model.get_inputs()[0].name: x.cpu().numpy()}
outs = onnx_model.run(None, inputs)
print(outs[0])

资料

https://blog.csdn.net/weixin_43198122/article/details/124452172