Refactor TensorRT engine export

[zhiqwang]

We provide a utilization tool export_tensorrt_engine for exporting TensorRT engines.

How to Export

import torch
from yolort.runtime.trt_helper import export_tensorrt_engine
from yolort.v5 import attempt_download

# Define some parameters
batch_size = 1
img_size = 640
score_thresh = 0.35
nms_thresh = 0.45
detections_per_img = 100
precision = "fp32"  # Currently only supports fp32

# 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"

input_sample = torch.rand(batch_size, 3, img_size, img_size)

export_tensorrt_engine(
    model_path,
    score_thresh=score_thresh,
    nms_thresh=nms_thresh,
    onnx_path=onnx_path,
    engine_path=engine_path,
    input_sample=input_sample,
    detections_per_img=detections_per_img,
)

Inference Interface

from yolort.runtime import PredictorTRT

# Load the exported TensorRT engine
engine_path = 'yolov5n6.engine'
size_divisible = 64  # for pre-processing
device = torch.device('cuda')
y_runtime = PredictorTRT(engine_path, device=device, size_divisible=size_divisible)

# Perform inference on an image file
predictions = y_runtime.predict('bus.jpg')

[CLAassistant]

All committers have signed the CLA.

[codecov]

Codecov Report

Merging #312 (8b5c98a) into main (bfc5d13) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##             main     #312   +/-   ##
=======================================
  Coverage   94.93%   94.94%           
=======================================
  Files          11       10    -1     
  Lines         731      732    +1     
=======================================
+ Hits          694      695    +1     
  Misses         37       37           

Flag	Coverage Δ
unittests	94.94% <100.00%> (+<0.01%)	⬆️

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

Impacted Files	Coverage Δ
test/test_relaying.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 bfc5d13...8b5c98a. Read the comment docs.