torch模型导出onnx

记录一下torch模型转onnx，onnx的优点是速度快，并且文件小，并且可支持较多的引擎，本文简单介绍如何导出：

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

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)

如此即可导出模型，但如果我们的输入模型参数是不断变化的呢，例如第三个维度的长度不是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])

如此即可。over