After training my own dataset, the labels of pt model inference and engine model inference are inconsistent.

[arkerman]

Search before asking

• [X] I have searched the YOLOv5 issues and discussions and found no similar questions.

Question

I set three class names : cls_001, cls_002, cls_003. And used command line below to train my dataset. python segment/train.py --data data/mydata.yaml --weights yolov5m-seg.pt --img 640 --cfg models/segment/yolov5m-seg.yaml --epochs 1000 Export engine. python export.py --weights runs/train-seg/exp13/weights/best.pt --include engine --device 0 --data data/mydata.yaml

Used pt model to infer. python segment/predict.py --weights runs/train-seg/exp13/weights/best.pt --source E:\data\mydata\test\images When i use pt model to predict, i get the correct labels.

Use engine to infer. python segment/predict.py --weights runs/train-seg/exp13/weights/best.engine --source E:\data\mydata\test\images When i use engine model to predict, i get the the label of "person", "bicycle", "car" on test images.

This confused me.

Additional

No response

[glenn-jocher]

@arkerman hello! Thank you for providing detailed information about your issue. It seems like there might be a discrepancy in how the class labels are being handled between the .pt model and the .engine model.

One possible cause could be related to the class labels not being correctly integrated or recognized in the conversion process to the TensorRT engine format. Here are a couple of steps you can take to troubleshoot this issue:

1. Verify Class Labels in YAML: Ensure that the mydata.yaml file used during the training and the export process has the correct class labels listed and in the same order.

2. Check Export Logs: Review the logs generated during the export process for any warnings or errors related to class labels or the conversion process.

3. Re-export the Model: Try re-exporting the .pt model to the .engine format, ensuring that all command line parameters are correctly specified and consistent with the training phase.

If these steps do not resolve the issue, please provide the contents of your mydata.yaml and any relevant logs from the export process, so we can help diagnose the problem further.

Thank you for using YOLOv5, and we appreciate your contribution to making the tool better! 🚀

[arkerman]

@glenn-jocher Hi! Really appreciate your suggestions! I debuged the predict.py code, and found it what cause the question. I forgot to add the argument --data. I used the command line below, and the everything is ok now. python segment/predict.py --weights runs/train-seg/exp13/weights/best.engine --source E:\data\mydata\test\images --data mydata.yaml Thanks again!

[glenn-jocher]

@arkerman That's great to hear! I'm glad you were able to resolve the issue by including the --data argument. If you have any more questions or run into further issues, feel free to reach out. Happy detecting! 🚀

[github-actions[bot]]

👋 Hello there! We wanted to give you a friendly reminder that this issue has not had any recent activity and may be closed soon, but don't worry - you can always reopen it if needed. If you still have any questions or concerns, please feel free to let us know how we can help.

For additional resources and information, please see the links below:

• Docs: https://docs.ultralytics.com
• HUB: https://hub.ultralytics.com
• Community: https://community.ultralytics.com

Feel free to inform us of any other issues you discover or feature requests that come to mind in the future. Pull Requests (PRs) are also always welcomed!

Thank you for your contributions to YOLO 🚀 and Vision AI ⭐