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

import torch
import onnxsim
import onnx
def export_to_onnx(model, output_path, input_shape, input_name, output_names):
    dummy_input = torch.rand(1, *input_shape)
    model.eval()
    temp_dict = dict()
    temp_onnx_path = output_path.replace('.onnx', '_temp.onnx')
    torch.onnx.export(model, 						# pytorch 模型
                    (dummy_input, 'ALL'),  			# 可以输入 tuple 
                    temp_onnx_path, 				# 输出 onnx 模型路径
                    verbose=False, 					# 聒噪
                    opset_version=11,				# onnx 版本
                    export_params=True, 			# 一个指示是否导出模型参数（权重）以及模型架构的标志。
                    do_constant_folding=True,   	# 一个指示是否在导出过程中折叠常量节点的标志
                    input_names=[input_name],		# 输入节点名称列表（可选）
                    output_names=output_names		# 输出节点名称列表（可选）
                    )
    input_data = {'image': dummy_input.cpu().numpy()}
    model_sim, flag = onnxsim.simplify(temp_onnx_path, input_data=input_data) # 简化 onnx 
    if flag:
        onnx.save(model_sim, output_path)
        print(f"simplify onnx model successfully !")
    else:
        print(f"simplify onnx model failed !!!")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

...
input_name = 'image'						# 输入节点名称
dynamic_axes = dict()
dynamic_axes[input_name] = [0, 2, 3]		# 输入节点开启动态输入的轴
        output_name = 'output'						# 输出节点名称
dynamic_axes[output_name] = [2, 3]			# 输出节点动态轴
torch.onnx.export(self, 
                (dummy_input, 'ALL'), 
                temp_onnx_path, 
                verbose=False, 
                opset_version=11,
                export_params=True, 
                do_constant_folding=True,
                input_names=[input_name],
                output_names=output_names,
                dynamic_axes=dynamic_axes	# 添加参数
                )

1
2
3
4
5
6
7
8

# onnx check
onnx_model = onnx.load(onnx_model_path)
try:
    onnx.checker.check_model(onnx_model)
except onnx.checker.ValidationError as e:
    print('The model is invalid: %s' % e)
else:
    print('The model is valid!')

1
2

pip install onnxruntime       # CPU build
pip install onnxruntime-gpu   # GPU build 

1
2
3
4
5
6

import onnxruntime
session = onnxruntime.InferenceSession("path to model")
session.get_modelmeta()
results = session.run(["output1", "output2"], {"input1": indata1, "input2": indata2})
results = session.run([], {"input1": indata1, "input2": indata2}) 

1
2
3
4
5

import onnxruntime
print(onnxruntime.get_available_providers())
# >>> ['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'AzureExecutionProvider', 'CPUExecutionProvider']

1
2
3
4
5

# 检查是否有'CUDAExecutionProvider'
if 'CUDAExecutionProvider' in onnxruntime.get_available_providers():
    print('ONNX Runtime GPU is available.')
else:
    print('ONNX Runtime GPU is not available.')

1
2
3
4
5
6
7

import onnxruntime as ort
# 加载ONNX模型
session = ort.InferenceSession('your_model.onnx', providers=['CUDAExecutionProvider'])
# 准备输入数据
inputs = {session.get_inputs()[0].name: your_input_data}
# 执行模型
outputs = session.run(None, inputs)