Fix TensorRT Python examples

[zhiqwang]

• Fix the TensorRT Python Notebook, check https://github.com/zhiqwang/yolov5-rt-stack/blob/main/notebooks/onnx-graphsurgeon-inference-tensorrt.ipynb for more details.
• set the default shape in ONNX TensorRT to (1, 3, 640, 640).

Or use the following Python scripts:

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 #283 (8d13c11) into main (3485ea1) will not change coverage.
The diff coverage is n/a.

@@           Coverage Diff           @@
##             main     #283   +/-   ##
=======================================
  Coverage   94.01%   94.01%           
=======================================
  Files          11       11           
  Lines         718      718           
=======================================
  Hits          675      675           
  Misses         43       43           

Flag	Coverage Δ
unittests	94.01% <ø> (ø)

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

---

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 3485ea1...8d13c11. Read the comment docs.