multi-class detection vs few-class detection ?

[hoangkhoiLE]

❔Question

Dear @glenn-jocher,
I want to do detection on person and bag object which are two class from COCO dataset. So, if i want to get a model better performance in only person and bag, do i need to retrain a model with label filtered between that two class and remove label of other classes ? Can you give me an advise in this situation ? Thank in advance !

Best regard !

Additional context

[glenn-jocher]

@hoangkhoiLE yes the fewer classes the model trains on the better it will perform at each class naturally. I've created PR #5172 to allow for class filtering by manually updating datasets.py.

To train on person and bag classes only set this variable in datasets.py accordingly:

        include_class = [0, 26]  # filter labels to include only these classes (optional)

[glenn-jocher]

@hoangkhoiLE good news 😃! Your original issue may now be fixed ✅ in PR #5172 which is now merged. To receive this update:

• Git – git pull from within your yolov5/ directory or git clone https://github.com/ultralytics/yolov5 again
• PyTorch Hub – Force-reload with model = torch.hub.load('ultralytics/yolov5', 'yolov5s', force_reload=True)
• Notebooks – View updated notebooks
• Docker – sudo docker pull ultralytics/yolov5:latest to update your image

Thank you for spotting this issue and informing us of the problem. Please let us know if this update resolves the issue for you, and feel free to inform us of any other issues you discover or feature requests that come to mind. Happy trainings with YOLOv5 🚀!

[MattVAD]

Hello @glenn-jocher, thank you for fixing this ! I'll try it soon. Just one more question, in your opinion/experience, if I want to train a model on only 2 classes (with the example of the amazing hoangkhoiLE, person & bag), is it better to use the pretrained weights or train the model from scratch ?

[glenn-jocher]

@MattVAD for most use cases (smaller datasets) starting from pretrained models will deliver better results. If you have large (COCO-size) dataset, then training from scratch may also produce good results.

I would train both ways if I were you to compare.