0

我使用这个 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 插件。

4

2 回答 2

0

我过去曾通过两种方式为 yoloV3 完成此操作：（两者都适用于 yolov4 和 yolov3-tiny）：

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

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

于 2020-12-03T14:27:04.930 回答
0

感谢您的帮助@joostblack。

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

于 2020-12-29T18:33:34.027 回答