-
Notifications
You must be signed in to change notification settings - Fork 0
/
pth2onnx.py
125 lines (92 loc) · 4.88 KB
/
pth2onnx.py
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
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
# -*- coding: utf-8 -*-
"""
@Time : 2024/8/11 16:48
@File : pth2onnx.py
@Author : zj
@Description:
Usage: Pytorch to ONNX:
$ python3 pth2onnx.py ./crnn_tiny-plate-b512-e100.pth ./runs/crnn_tiny-plate.onnx
$ python3 pth2onnx.py ./crnn_tiny-emnist-b512-e100.pth ./runs/crnn_tiny-emnist.onnx
"""
import argparse
import os.path
import numpy as np
import onnx
import onnxruntime
import torch.onnx
import torch.nn as nn
from utils.general import load_crnn
from utils.dataset.emnist import DIGITS_CHARS
from utils.dataset.plate import PLATE_CHARS
def parse_opt():
parser = argparse.ArgumentParser(description="Pytorch to ONNX")
parser.add_argument("pretrained", metavar="MODEL", type=str, default=None, help="Pytorch Pretrained Model Path")
parser.add_argument("save", metavar="SAVE", type=str, default=None, help="Saving ONNX Path")
parser.add_argument('--use-lstm', action='store_true', help='use nn.LSTM instead of nn.GRU')
parser.add_argument('--not-tiny', action='store_true', help='Use this flag to specify non-tiny mode')
args = parser.parse_args()
print(f"args: {args}")
return args
def check_onnx(onnx_path='pytorch.onnx'):
onnx_model = onnx.load(onnx_path)
onnx.checker.check_model(onnx_model)
def check_output(x, torch_out, onnx_path='pytorch.onnx'):
# See https://blog.csdn.net/zunzunle/article/details/130087922
print("Supported onnxruntime version: ", onnxruntime.__version__)
print("Supported Opset versions: ", onnxruntime.get_available_providers())
# ValueError: This ORT build has ['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider'] enabled. \
# Since ORT 1.9, you are required to explicitly set the providers parameter when instantiating InferenceSession.
# For example, onnxruntime.InferenceSession(..., providers=['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider'], ...)
ort_session = onnxruntime.InferenceSession(onnx_path, providers=['CPUExecutionProvider'])
print("Onnx info:")
print(f" input: {ort_session.get_inputs()[0]}")
print(f" output: {ort_session.get_outputs()[0]}")
def to_numpy(tensor):
return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
# compute ONNX Runtime output prediction
ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x)}
ort_outs = ort_session.run(None, ort_inputs)
print(x.shape, ort_outs[0].shape)
# compare ONNX Runtime and PyTorch results
np.testing.assert_allclose(to_numpy(torch_out), ort_outs[0], rtol=1e-03, atol=1e-05)
print("Exported model has been tested with ONNXRuntime, and the result looks good!")
def export_to_onnx(torch_model, shape=None, onnx_path="pytorch.onnx", is_dynamic=False):
assert isinstance(torch_model, nn.Module)
# Input to the model
x = torch.randn(shape, requires_grad=True)
torch_out = torch_model(x)
# Export the model
if is_dynamic:
# variable length axes
dynamic_axes = {'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}
else:
dynamic_axes = None
torch.onnx.export(torch_model, # model being run
x, # model input (or a tuple for multiple inputs)
onnx_path, # 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=14, # 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=dynamic_axes
)
check_onnx(onnx_path=onnx_path)
check_output(x, torch_out, onnx_path=onnx_path)
def main(args):
# UserWarning: Exporting a model to ONNX with a batch_size other than 1, with a variable length with GRU can cause an error when running the ONNX model with a different batch size.
# Make sure to save the model with a batch size of 1, or define the initial states (h0/c0) as inputs of the model.
if 'plate' in os.path.basename(args.pretrained):
shape = (1, 3, 48, 168)
num_classes = len(PLATE_CHARS)
else:
shape = (1, 1, 32, 160)
num_classes = len(DIGITS_CHARS)
model, device = load_crnn(pretrained=args.pretrained, shape=shape, num_classes=num_classes,
not_tiny=args.not_tiny, use_lstm=args.use_lstm)
onnx_path = args.save
export_to_onnx(model, shape=shape, onnx_path=onnx_path, is_dynamic=False)
print(f"Save to {onnx_path}")
if __name__ == '__main__':
args = parse_opt()
main(args)