Yolov3 到 Tensorrt：tf-keras Lambda 层的自定义插件

Question

我使用这个 repo 在 Tensorflow 2.0 中训练了一个 yolov3-tiny 模型：https://github.com/zzh8829/yolov3-tf2

在推理时，该模型使用两个包裹在 tf-keras lambda 层中的函数进行后处理，它们是：

• yolo_boxes : 根据模型输出的偏移量计算实际的盒子坐标
• yolo_nms : 使用 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]))

我已经创建了这个推理模型的冻结 pb，并将其转换为 ONNX。但我无法弄清楚如何进行。 如何为 yolo_boxes 创建 python Tensorrt 插件？我在网上找不到任何 Lambda 层插件的材料，并且如果没有 yolo_boxes 插件，我无法测试 tensorrts 自定义 NMS 插件。

Answer 1

我过去曾通过两种方式为 yoloV3 做到这一点：（两者都适用于 yolov4 和 yolov3-tiny）：

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

它们都先转换为 ONNX，然后再转换为 TensorRT。 对于第二个链接，您将需要 Pytorch。

请注意，需要正确版本的 ONNX 和 TensorRT 才能完成这项工作。旧版本的 ONNX 没有正确的 opset 可以工作。但是这些信息都可以在这两个链接上找到。

Answer 2

感谢您的帮助@joostblack。

参考您的资源和其他一些资源，我能够使用 BatchedNMSDynamic_TRT 插件来解决我的问题。我只是将 combine_non_max_suppression 函数（yolo_nms 中的那个）的输入作为模型的输出，并使用 graphsurgeon 将 batchedNMSDynamic_TRT 插件附加为节点，将模型的输出作为插件的输入。