层的自定义插件Yolov3到坦索尔特

Question

我在Tensorflow 2.0中训练了一个yolov3 3-微型模型，使用这个回购程序：https://github.com/zzh8829/yolov3-tf2

根据推断，该模型使用封装在tf-keras lambda层中的两个函数进行后处理，这些函数是：

tf.image.combined_non_max_suppression

• yolo_boxes :从模型
• yolo_nms输出的偏移量中计算实际的盒坐标:使用
• 进行非最大抑制

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插件？我无法在网上找到任何材料的Lambda层插件，也不能测试张力自定义NMS插件，没有yolo_boxes插件的第一。

Answer 1

过去，我用两种方式为yoloV3做过这件事：(这两种方法都适用于yolov4和yolov4 3-微)：

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

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

请注意，要使此工作正常，需要正确版本的ONNX和TensorRT。旧版本的ONNX没有正确的选项来工作。但是这些信息都可以在这两个链接上找到。

Answer 2

谢谢你的帮助@joostblack。

参考您的资源和其他一些资源，我能够使用BatchedNMSDynamic_TRT插件来解决我的问题。我只是将输入到combined_non_max_suppression函数(yolo_nms中的函数)作为模型的输出，并使用图形外科医生将batchedNMSDynamic_TRT插件作为节点附加，将模型的输出作为插件的输入。