jetson-inference icon indicating copy to clipboard operation
jetson-inference copied to clipboard

How to add custom classes while preserving the classes in original model

Open AgemennoN opened this issue 2 years ago • 7 comments

Hi Dusty,

I collected my own data and trained a custom model by following steps in here (https://github.com/dusty-nv/jetson-inference/blob/master/docs/pytorch-collect-detection.md). It works fine. But now I want to add my new classes to ssd_mobilenet-v1. I mean if the ssd_mobilenet has 90 classes I want to make it 92 classes(90 original classes + 2 my own classes) and I couldn't find how. Can you help me please?

AgemennoN avatar Feb 10 '22 13:02 AgemennoN

Hi @AgemennoN, that 90-class ssd-mobilenet-coco model came from TensorFlow, whereas the new models are trained in PyTorch, so that model in particular isn't capable of being extended (unless you were to figure it out in TensorFlow which I don't recommend for various reasons)

But lets say you wanted to re-train a PyTorch ssd-mobilenet that was initially trained on the Pascal VOC dataset. If you made a new dataset in Pascal VOC format with your new data, then trained the model again while keeping the old+new classes in labels.txt, that could in theory work.

However I think the most straightforward way given the train_ssd.py scripts as they exist today, would be to just make one dataset with all your data and train the model that way. I'm not an expert on training so I'm sure there are better ways, however I just try to use the training scripts that are already working.

dusty-nv avatar Feb 10 '22 14:02 dusty-nv

Hi, @dusty-nv

Thanks for your answer.

I wanted to try finding a model trained with Pascal VOC dataset. But I have a few more questions

  1. I will use the model as "--pretrained-ssd=mobilenet-v1-ssd-voc" while executing train_ssd.py right? But what about the "--net" flag. I couldn't understand its purpose.

  2. I couldn't manage to find a ssd-mobilenet trained with Pascal VOC, I believe it doesn't have to be mobilenet so I tried to use FaceBoxesProd.pth from here "https://github.com/zisianw/FaceBoxes.PyTorch". But I guess being a ssd model with trained VOC dataset isn't enough.

  3. In your third option you meant, use a dataset that includes all data including all the 92 classes right?

Thank you in advance

AgemennoN avatar Feb 11 '22 13:02 AgemennoN

  1. I will use the model as "--pretrained-ssd=mobilenet-v1-ssd-voc" while executing train_ssd.py right? But what about the "--net" flag. I couldn't understand its purpose.

There's a pretrained PyTorch checkpoint that comes with the repo that was already trained on Pascal VOC. It's actually the default setting to the --pretrained-ssd argument: https://github.com/dusty-nv/pytorch-ssd/blob/8ed842a408f8c4a8812f430cf8063e0b93a56803/train_ssd.py#L51

The --net flag changes the network architecture, not the model. Keep it as --net=mb1-ssd (the default)

dusty-nv avatar Feb 11 '22 16:02 dusty-nv

So, if I create a label.txt file with a few class from pretrained model with my new classes, then I can train a new model only with the data of my new classes, in theory it should detect the objects in the label.txt file. Is that correct?

AgemennoN avatar Feb 14 '22 06:02 AgemennoN

I'm not entirely sure, but I think the new labels.txt would need to have all the old classes (in the same order), in addition to your new classes. If you changed the structure/order of the old classes, you would need to re-train those too.

dusty-nv avatar Feb 14 '22 18:02 dusty-nv

@AgemennoN You manage to use the previous pretrained data with new data? Thanks

miguemorales avatar Jul 20 '22 08:07 miguemorales

@AgemennoN You manage to use the previous pretrained data with new data? Thanks

Unfortunately I couldn't and I gave up (It was the last week of my internship and Jetson belonged to the company)

AgemennoN avatar Jul 20 '22 12:07 AgemennoN