Inference with TensorRT

[Naofel-eal]

Hello, I have some issues trying to infer an image with YOLO-NAS-S model converted to .engine. Has anyone managed to do this ?

I used this script to export the custom model to .onnx:

from super_gradients.common.object_names import Models
from super_gradients.training import models
import torch

model = models.get(Models.YOLO_NAS_S, checkpoint_path='checkpoints/model.pth', num_classes=2)

model.eval()
model.prep_model_for_conversion(input_size=[1, 3, 640, 640])

dummy_input = torch.randn(1, 3, 640, 640) 

torch.onnx.export(model, dummy_input,  "checkpoints/model.onnx")

And this command to convert the .onnx to .engine:

tensorrt/bin/trtexec --onnx=model.onnx --workspace=2048 --avgRuns=100 --duration=15 --int8 --fp16 --saveEngine=model.engine

But all the scripts I tested to infer an image failed. Does anyone have a working python program that infers an image ?

[dagshub[bot]]

Join the discussion on DagsHub!

[BloodAxe]

Summon TRT guru @NatanBagrov

[pranoyr]

Try with this repo https://github.com/Linaom1214/TensorRT-For-YOLO-Series

ONNX --> TensorRT python export.py -o yolo_nas_m.onnx -e yolonas.trt

For inference use this --> python trt.py -e yolo_nas.trt

Edit utils/util.py , Add this after line 107, you can see the results.

boxes = torch.from_numpy(data[1]).view(1,8400,4)
scores = torch.from_numpy(data[0]).view(1,8400,80)
print(f"boxes: {boxes.shape}")
print(f"outputs: {scores.shape}")

[NatanBagrov]

Hi @Naofel-eal , could you share more details about the error messages you're getting?

[Naofel-eal]

Hi @NatanBagrov, For example, I tested the repo mentioned by @pranoyr (thank you for your help !) but I got the following error during the inference:

\utils\utils.py", line 109, in inference
    num, final_boxes, final_scores, final_cls_inds = data
ValueError: not enough values to unpack (expected 4, got 2)

I also tested this stackoverflow answer but I don't know how to get the bounding box from the output. The result has a length of 2, result[0].shape = (2, 8400) and result[1].shape = (2, 16800) and I don't know what it means. Answer: https://stackoverflow.com/questions/59280745/inference-with-tensorrt-engine-file-on-python/67492525#67492525

[NatanBagrov]

Hmmm... that's weird. From what I see here: https://github.com/Deci-AI/super-gradients/blob/master/src/super_gradients/training/models/detection_models/yolo_nas/dfl_heads.py#L225-L228 I would expect 2 tensors, one with a shape of (4, 8400) and the other with a shape of (80, 8400), or something like that. Maybe @BloodAxe or @Louis-Dupont can help

[Naofel-eal]

Maybe I made a mistake when exporting to .onnx or .trt, I am a beginner. Is the export script in the first message correct ?

[Naofel-eal]

I tested this repo https://github.com/Linaom1214/TensorRT-For-YOLO-Series/blob/main/trt.py. The inference works very well with a yolov8.trt model but not for the yolonas.trt model. I think the conversion for yolonas from .pth to .trt doesn't works well. Is there any offficial way to do that ?

[Louis-Dupont]

We just added model.export API which strongly simplifies that flow, it is available in master and will be released in 3.2.0 https://github.com/Deci-AI/super-gradients/blob/master/documentation/source/models_export.md

This will soon be added to the documentation interface