Retrain model using less categories without any loss

AI & Data Science Deep Learning (Training & Inference) TensorRT
1

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

2

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!

4

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.

5

Hi @Rik ,

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

This page talks about the issues related to TensorRT.

Thanks