Hi, i’m trying to replicate the TensorRT implementations for the models shown here https://github.com/AastaNV/TRT_object_detection
but using the tf2onnx instead of using the uff parser.
Is there some guidance for how to do this? I am mostly using the code discussed here as guide.
My specific error is from the grid anchor plugin,
[06/27/2021-17:10:52] [F] [TRT] Assertion failed: numExpectedLayers == numLayers
gridAnchorPlugin.cpp:454
I’m using the same gridAnchor parameters as the uff examples. Here’s the code. The input onnx graph was created from from a SSD_Inception_V2 frozen_graph.pb file using tf2onnx. This pb file can be converted can be converted to uff and run with tensorRT as shown in here GitHub - AastaNV/TRT_object_detection: Python sample for referencing object detection model with TensorRT
import onnx_graphsurgeon as gs
import argparse
import onnx
import numpy as np
####
#### Model parameters
####
num_classes = 9
init_inputs = ["image_tensor:0"]
init_outputs = ["detection_boxes:0",
"detection_scores:0",
"detection_classes:0",
"num_detections:0",]
def generate_prior_var_shapes(batch_size):
prior_var_shapes = [[batch_size,2,3*19*19*4,1],
[batch_size,2,6*10*10*4,1],
[batch_size,2,6*5*5*4,1],
[batch_size,2,6*3*3*4,1],
[batch_size,2,6*2*2*4,1],
[batch_size,2,6*1*1*4,1],
]
return prior_var_shapes
prior_box_attrs = {
'numLayers':6,
'minSize':0.2,
'maxSize':0.95,
'aspectRatios':[1.0, 2.0, 0.5, 3.0, 0.33],
'variance':[0.1,0.1,0.2,0.2],
'featureMapShapes':[19, 10, 5, 3, 2, 1],
}
NMS_attrs = {
'shareLocation':1,
'varianceEncodedInTarget':0,
'backgroundLabelId':0,
'confidenceThreshold':1e-8,
'nmsThreshold':0.6,
'topK':100,
'keepTopK':100,
'numClasses':9,
'inputOrder':[0, 2, 1],
'confSigmoid':1,
'isNormalized':1
}
def create_and_add_plugin_node(graph,input_shape=None):
if input_shape is None:
batch_size = graph.inputs[0].shape[0]
input_h = graph.inputs[0].shape[1]
input_w = graph.inputs[0].shape[2]
channels = graph.inputs[0].shape[3]
else:
batch_size,input_h,input_w,channels = input_shape
nodes = graph.nodes
tensors = graph.tensors()
prior_var_shapes = generate_prior_var_shapes(batch_size)
grid_anchor_outs = []
total_boxes = 0
for idx,data in enumerate(zip(prior_box_attrs['featureMapShapes'],prior_var_shapes)):
nprior,var_shape = data
total_boxes+=var_shape[2]
grid_anchor_outs.append(gs.Variable(name="GridAnchor_{}_out".format(idx)).to_variable(dtype=np.int32,
shape=var_shape))
box_node = gs.Node(name= 'GridAnchor',
op="GridAnchor_TRT",
attrs=prior_box_attrs,
outputs=grid_anchor_outs
)
concat_priorbox_out = gs.Variable(name="concat_priorbox_out").to_variable(dtype=np.float32,
shape=[batch_size,2,total_boxes,1])
concat_priorbox = gs.Node(name="concat_priorbox", op="ConcatV2",
attrs={'axis':2},
inputs=grid_anchor_outs,
outputs=[concat_priorbox_out])
graph.nodes.append(concat_priorbox)
graph.nodes.append(box_node)
###
### Below has more to do with NMS, not directly related to gridAnchor error
###
original_concat_box_loc = tensors["concat:0"]
original_concat_box_conf = tensors["concat_1:0"]
keepTopK = NMS_attrs['keepTopK']
NMS_out_0 = gs.Variable(name="NMS_out_0").to_variable(dtype=np.float32,
shape=[batch_size,1,keepTopK,7])
NMS_out_1 = gs.Variable(name="NMS_out_1").to_variable(dtype=np.int32,
shape=[batch_size,1,1,1])
nms_node = gs.Node(name="NMS",
op="NMS_TRT",
attrs=NMS_attrs,
inputs=[original_concat_box_loc, concat_priorbox_out,original_concat_box_conf],
outputs=[NMS_out_0,NMS_out_1])
graph.nodes.append(nms_node)
graph.outputs = [NMS_out_0]
return graph.cleanup().toposort()
def main():
parser = argparse.ArgumentParser(description="Add NMS")
parser.add_argument("-onnx_path", "--onnx_path",
help="Path to the ONNX model generated by export_model.py",
)
parser.add_argument("--input_shape", "--input_shape", default=None,type=int,nargs='+',
help="input shape as list (in NHWC order) e.g. '-1 300 300 3'",
)
args, _ = parser.parse_known_args()
graph = gs.import_onnx(onnx.load(args.onnx_path))
graph = create_and_add_plugin_node(graph,input_shape=args.input_shape)
onnx.save(gs.export_onnx(graph), args.onnx_path + ".nms.onnx")
if __name__ == '__main__':
main()
Environment
TensorRT Version: 8.0
GPU Type: 1070
Nvidia Driver Version: 460.73.01
CUDA Version: 11.2
CUDNN Version: 8
Operating System + Version: Ubuntu 18.04.5 LTS
Python Version (if applicable): 3.7
TensorFlow Version (if applicable): 2.5