-
Notifications
You must be signed in to change notification settings - Fork 1
/
diyierpi_onnx2trt_multibatch.py
148 lines (122 loc) · 5.6 KB
/
diyierpi_onnx2trt_multibatch.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
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
# This sample uses an ONNX ResNet50 Model to create a TensorRT Inference Engine
import random
from PIL import Image
import numpy as np
import cv2
import pycuda.driver as cuda
# This import causes pycuda to automatically manage CUDA context creation and cleanup.
import pycuda.autoinit
import tensorrt as trt
import sys, os
sys.path.insert(1, os.path.join(sys.path[0], ".."))
import common
engine_file_path="diyierpi_branch_pt.trt"
model_path = "diyierpi_branch_pt.onnx"
make_num = 43
class_labels = "diyierpi_make_balance.txt"
class ModelData(object):
MODEL_PATH = model_path
INPUT_SHAPE = (3, 448, 448)
BATCH_SIZE = 16
# We can convert TensorRT data types to numpy types with trt.nptype()
DTYPE = trt.float32
# You can set the logger severity higher to suppress messages (or lower to display more messages).
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
# The Onnx path is used for Onnx models.
def build_engine_onnx(model_file):
with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
builder.max_workspace_size = common.GiB(1)
builder.max_batch_size = ModelData.BATCH_SIZE
# Load the Onnx model and parse it in order to populate the TensorRT network.
with open(model_file, 'rb') as model:
parser.parse(model.read())
# print(parser.parse(model.read()))
# print(parser.num_errors)
# print(parser.get_error(0))
# last_layer0 = network.get_layer(network.num_layers -11) #output 6
# last_layer1 = network.get_layer(network.num_layers -6) #185 463
# last_layer2 = network.get_layer(network.num_layers -1) #187 8
# print("-------")
# print(last_layer.get_output(0).name)
# print(last_layer.get_output(0).shape)
# print("-------")
# network.mark_output(last_layer0.get_output(0))
# network.mark_output(last_layer1.get_output(0))
# network.mark_output(last_layer2.get_output(0))
engine = builder.build_cuda_engine(network)
f = open(engine_file_path, "wb")
f.write(engine.serialize())
return engine
def load_normalized_test_case(test_image_txt):
# Converts the input image to a CHW Numpy array
def normalize_image(image):
# Resize, antialias and transpose the image to CHW.
c, h, w = ModelData.INPUT_SHAPE
image_arr = np.asarray(image.resize((w, h), Image.ANTIALIAS)).transpose([2, 0, 1]).astype(trt.nptype(ModelData.DTYPE)).ravel()
# This particular ResNet50 model requires some preprocessing, specifically, mean normalization.
return (image_arr / 255.0 - 0.45) / 0.225
f_r = open(test_image_txt,'r')
lines = f_r.readlines()
images=[]
imgname=[]
for line in lines:
img_path = line.strip()
imgname.append(img_path)
images.append(normalize_image(Image.open(img_path)))
return images,imgname
def softmax(x):
x = x - np.max(x)
exp_x = np.exp(x)
softmax_x = exp_x / np.sum(exp_x)
return softmax_x
cd={"黑":0,"蓝":1,"青":2,"灰":3,"绿":4,"红":5,"白":6,"黄":7 }
td={"sedan":0,"suv":1,"mini_bus":2,"bus":3,"truck":4,"van":5}
new_cd = {v : k for k, v in cd.items()}
new_td = {v : k for k, v in td.items()}
def main():
# Set the data path to the directory that contains the trained models and test images for inference.
# data_files = common.find_sample_data(description="Runs a ResNet50 network with a TensorRT inference engine.", subfolder="", find_files=[ModelData.MODEL_PATH, class_labels])
# Get test images, models and labels.
# print(len(data_files))
# test_images = data_files[0:3]
# onnx_model_file, labels_file = data_files[3:]
# test_images = data_files[0]
# print(test_images)
# onnx_model_file, labels_file = data_files[1:]
onnx_model_file = ModelData.MODEL_PATH
labels_file = class_labels
# print(onnx_model_file)
labels = open(labels_file, 'r').read().split('\n')
engine = build_engine_onnx(onnx_model_file)
# h_input, d_input, h_output0, d_output0, h_output1, d_output1, h_output2, d_output2, stream = allocate_buffers(engine)
inputs, outputs, bindings, stream = common.allocate_buffers(engine)
context = engine.create_execution_context()
test_image_txt = 'test.txt'
images,imgname = load_normalized_test_case(test_image_txt)
inputs[0].host = np.array(images) #
# test_case = load_normalized_test_case(test_images, h_input)
# do_inference(context, h_input, d_input, h_output0, d_output0,h_output1, d_output1,h_output2, d_output2, stream)
trt_outputs = common.do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream, batch_size=ModelData.BATCH_SIZE)
h_output0 = trt_outputs[0].reshape([-1, 6])
# h_output1 = trt_outputs[1].reshape([-1, 463])
h_output1 = trt_outputs[1].reshape([-1, make_num])
h_output2 = trt_outputs[2].reshape([-1, 8])
for i in range(0,len(images)):
# print(len(h_output0))
# print(len(h_output1))
# print(len(h_output2))
print("----------------")
print(imgname[i])
prob0 = np.max(softmax(h_output0[i]))
pred0 = new_td[np.argmax(softmax(h_output0[i]))]
print("type: ", pred0,prob0)
prob1 = np.max(softmax(h_output1[i]))
pred1 = labels[np.argmax(softmax(h_output1[i]))]
# print("makemodel: ", pred1,prob1)
print("make: ", pred1,prob1)
prob2 = np.max(softmax(h_output2[i]))
pred2 = new_cd[np.argmax(softmax(h_output2[i]))]
print("color: ", pred2,prob2)
print("----------------\n")
if __name__ == '__main__':
main()