Yolov3 to Tensorrt : Custom Plugins for tf-keras Lambda Layers

Question

I have trained a yolov3-tiny model in Tensorflow 2.0 using this repo : https://github.com/zzh8829/yolov3-tf2

On inference, the model uses two functions wrapped in tf-keras lambda layers for postprocessing, these are :

• yolo_boxes : to calculate actual box coordinates from the offsets outputted by the model
• yolo_nms : do nonmax-suppression using tf.image.combined_non_max_suppression

boxes_0 = Lambda(lambda x: yolo_boxes(x, anchors[masks[0]], classes),name='yolo_boxes_0')(output_0)

boxes_1 = Lambda(lambda x: yolo_boxes(x, anchors[masks[1]], classes),name='yolo_boxes_1')(output_1)

outputs = Lambda(lambda x: yolo_nms(x, anchors, masks, classes),name='yolo_nms')((boxes_0[:3], boxes_1[:3]))

I have created a frozen pb of this inference model, and converted it to ONNX. But I cannot figure out how to proceed. How do I create a python Tensorrt plugin for yolo_boxes? I cannot find any material online for Lambda layer plugins, and cannot test tensorrts custom NMS plugin without the yolo_boxes plugin first.

Answer 1

I have done this in the past for yoloV3 in two ways: (Both also work for yolov4 and yolov3-tiny):

• https://github.com/jkjung-avt/tensorrt_demos
• https://github.com/Tianxiaomo/pytorch-YOLOv4

They both first convert to ONNX and then to TensorRT. For the second link you will need Pytorch.

Note that the right versions of ONNX and TensorRT are required to make this work. Old versions of ONNX do not have the right opset to work. But this information can all be found on those two links.

Answer 2

Thank you for your help @joostblack.

Referring to your resources and a few others, I was able to use BatchedNMSDynamic_TRT plugin to solve my problem. I simply made the input to the combined_non_max_suppression function (the one within yolo_nms) as the output of the model, and used graphsurgeon to append the batchedNMSDynamic_TRT plugin as a node, giving the outputs of the model as inputs to the plugin.