jetson-inference
jetson-inference copied to clipboard
How to add custom classes while preserving the classes in original model
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?
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.
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
-
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.
-
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.
-
In your third option you meant, use a dataset that includes all data including all the 92 classes right?
Thank you in advance
- 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)
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?
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.
@AgemennoN You manage to use the previous pretrained data with new data? Thanks
@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)