Retrain model using less categories without any loss

Description

Hello, i have a pretrained yolov5 model trained on coco dataset with 80categories
currently i would like to remove from the model 78 categories and use just two of them, what’s the best way to do it?
do i need to download the coco dataset, select just the two categories of interest and retrain the model from scratch?
do i need to download the coco dataset, select just the two categories of interest and retrain the model starting from the checkpoint that i currently have?

Is there any other way to reach this goal?

Environment

**Training **: google colab
**Training **: NVIDIA Tesla
deploy
**TensorRT **:8
GPU Type: Nvidia Jetson
CUDA Version: 10.2
Operating System: Ubuntu 1804
Python Version (if applicable): 3.6.9

Model

GitHub - ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite here you can find the model that i am using

Thanks guys

Hi,
Request you to share the ONNX model and the script if not shared already so that we can assist you better.
Alongside you can try few things:

  1. validating your model with the below snippet

check_model.py

import sys
import onnx
filename = yourONNXmodel
model = onnx.load(filename)
onnx.checker.check_model(model).
2) Try running your model with trtexec command.

In case you are still facing issue, request you to share the trtexec “”–verbose"" log for further debugging
Thanks!

hi,
does not matter the ONNX, i think that the problem is model intependent, i need to know how retrain a model with a subset of it’s categories.

Hi @Rik ,

I believe you can raise the concern on the respective github page.

This page talks about the issues related to TensorRT.

Thanks