Yolov3 到 Tensorrt：tf-keras Lambda 層的自定義插件

Question

我使用這個 repo 在 Tensorflow 2.0 中訓練了一個 yolov3-tiny model： https://github.com/zzh8829/yolov3-tf2

據推斷，model 使用兩個封裝在 tf-keras lambda 層中的函數進行后處理，它們是：

• yolo_boxes：根據 model 輸出的偏移量計算實際框坐標
• 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]))

我已經創建了這個推理 model 的凍結 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 插件來解決我的問題。 我只是將 combined_non_max_suppression function（yolo_nms 中的那個）作為 model 的 output 進行輸入，並將 graphsurgeon 用於 append 批處理的 NMSDynamic_TRT 插件作為節點，將 8957605 的輸入作為 8957616 插件的輸入