wrong output with segment/predict.py on tensorrt engine for custom dataset

Created by: GreatV

Search before asking

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

YOLOv5 Component

Detection, Export

Bug

here is the output of .engine model

 python segment/predict.py --source ~/train_dataset/1.png --weights runs/train-seg/exp/weights/best.engine

1.png

Here is the output of .onnx model

 python segment/predict.py --source ~/train_dataset/1.png --weights runs/train-seg/exp/weights/best.onnx

99240191

Environment

  • YOLOv5 🚀 v7.0-5-gbfa1f23 Python-3.9.15 torch-1.13.0 CUDA:0 (Tesla T4, 14961MiB)
  • Ubuntu 20.04
  • Python 3.9.15

Minimal Reproducible Example

python export.py  --weights runs/train-seg/exp/weights/best.pt --include engine --simplify --device 0

 python segment/predict.py --source ~/train_dataset/1.png --weights runs/train-seg/exp/weights/best.engine

Additional

TensorRT: input "images" with shape(1, 3, 640, 640) DataType.FLOAT
TensorRT: output "output0" with shape(1, 25200, 38) DataType.FLOAT
TensorRT: output "output1" with shape(1, 32, 160, 160) DataType.FLOAT
TensorRT: building FP32 engine as runs/train-seg/exp/weights/best.engine

Are you willing to submit a PR?

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