Tensorrt export fail

[dariogonle]

Search before asking

• I have searched the YOLOv5 issues and found no similar bug report.

YOLOv5 Component

Export

Bug

I've used the export.py script to convert yolov5n.pt to tensorrt.

python3 export.py --weights yolov5n.pt --include engine

But the conversion process fails.

It says that tensorrt is not installed, but it is:

Environment

No response

Minimal Reproducible Example

No response

Additional

No response

Are you willing to submit a PR?

• Yes I'd like to help by submitting a PR!

[dariogonle]

I think the problem is that for using tensorrt it is necessary to force to use the GPU and FP16. Therefore, the correct command should be

python export.py --include engine --device 0 --half --weights yolov5n.pt --data data.yaml

[glenn-jocher]

@dariogonle yes, TRT exports must be run on a GPU backend. Maybe we should have a special assert here for this.