Fix register_nms in YOLOGraphSurgeon

[zhiqwang]

Register the BatchedNMS_TRT plugin node at the end of output of the YOLOTRTModule, and now we can get an end-to-end ONNX graph except the pre-processing.

import os
import torch

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

assert torch.cuda.is_available()
device = torch.device('cuda')

from yolort.utils import get_image_from_url, read_image_to_tensor
from yolort.v5 import letterbox, scale_coords, attempt_download
from yolort.runtime import PredictorTRT
from yolort.runtime.trt_helper import EngineBuilder
from yolort.runtime.yolo_graphsurgeon import YOLOGraphSurgeon

# Define some parameters
img_size = 640
stride = 64
score_thresh = 0.35
iou_thresh = 0.45
detections_per_img = 100
half = False

# yolov5s6.pt is downloaded from 'https://github.com/ultralytics/yolov5/releases/download/v6.0/yolov5n6.pt'
model_path = "yolov5n6.pt"

checkpoint_path = attempt_download(model_path)
onnx_path = "yolov5n6.onnx"
engine_path = "yolov5n6.engine"

img_source = "https://huggingface.co/spaces/zhiqwang/assets/resolve/main/bus.jpg"
# img_source = "https://huggingface.co/spaces/zhiqwang/assets/resolve/main/zidane.jpg"
img_raw = get_image_from_url(img_source)

# Pre Processing
image = letterbox(img_raw, new_shape=(img_size, img_size), stride=stride)[0]
image = read_image_to_tensor(image)
image = image[None]
image = image.to(device)
image = image.contiguous()

# Export to ONNX models
yolo_gs = YOLOGraphSurgeon(model_path, version="r6.0", input_sample=image, enable_dynamic=False)
# Embed the `BatchedNMS_TRT` at the end of `LogitsDecoder`.
yolo_gs.register_nms(score_thresh=score_thresh, nms_thresh=iou_thresh, detections_per_img=detections_per_img)

yolo_gs.save(onnx_path)

# Build TensorRT Engine
engine_builder = EngineBuilder()
engine_builder.create_network(onnx_path)
engine_builder.create_engine(engine_path, precision="fp32")

# Inference on TensorRT
engine = PredictorTRT(engine_path, device)
engine.warmup(img_size=image.shape, half=half)

# Inferencing
detections = engine.run_on_image(image)

[CLAassistant]

All committers have signed the CLA.

[codecov]

Codecov Report

Merging #252 (ec6a4ec) into main (5105b37) will not change coverage.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##             main     #252   +/-   ##
=======================================
  Coverage   93.62%   93.62%           
=======================================
  Files          11       11           
  Lines         675      675           
=======================================
  Hits          632      632           
  Misses         43       43           

Flag	Coverage Δ
unittests	93.62% <100.00%> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files	Coverage Δ
test/test_runtime.py	100.00% <100.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 5105b37...ec6a4ec. Read the comment docs.