在使用torch自带的方法将训练好的权重文件转换成onnx格式时，输入的张量可以是静态的，也可以通过一些设置修改成动态的。

静态：在导出onnx时传入输入尺寸，在之后的预测过程中也必须使用相同尺寸的输入。

这种方法对于像分割检测这种任务不太友好，我们希望模型能够自适应输入尺寸，所以需要设置成动态输入形式。

设置方法也很方便，只需在导出时传入dynamic_axes，指定要动态输入的维度即可。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
import argparse
import os
import numpy as np
import time
import cv2

from modeling.deeplab import *
from dataloaders import custom_transforms as tr
from PIL import Image
from torchvision import transforms
from dataloaders.utils import  *
from torchvision.utils import make_grid, save_image

import onnxruntime

model = DeepLab(num_classes=3,
                    backbone='ghostnet',
                    output_stride=8,
                    sync_bn=None,
                    freeze_bn=False)
ckpt_path='model_best.pth.tar'
ckpt = torch.load(ckpt_path, map_location='cpu')
model.load_state_dict(ckpt['state_dict'])
model.eval()

x=torch.randn(1,3,224,224)

with torch.no_grad():
    print('tensor_in shape',x.shape)# 1x3x400x400
    output = model(x)

torch.onnx.export(model,x,'model.onnx',export_params=True,opset_version=11,do_constant_folding=True,input_names=['input'],output_names=['output'],dynamic_axes={"input":{2:'in_width',3:'in_height'}})


ort_session=onnxruntime.InferenceSession('model.onnx')

ort_inputs={ort_session.get_inputs()[0].name:x.numpy()}

ort_outs=ort_session.run(None,ort_inputs)

print(ort_outs)
print(output)

np.testing.assert_allclose(output.numpy(),ort_outs[0],rtol=1e-03,atol=1e-05)
print('successfully converted pth into onnx!')

上述代码输出一个model.onnx文件

在得到onnx文件后，可以使用onnxruntime进行推理：

参考：