Retrained SSD-Mobilenet model only has 1 class when running detectnet

Hi @dusty_nv, I did transfer learning on SSD-Mobilenet-v1 according your instructions at pytorch-ssd-transfer-learning-instructions, and when I run detectnet with this retrained model it only shows the new trained class (head) but doesn’t actually include the name (just calls it Class #1). And it can’t it identify the other 80 or more classes that this model was trained on?

Here’s a picture of what I’m seeing with the newly retrained model. I should see a box around my entire body labeled as “person” AND the bounding box around my head labeled “head”, and maybe it’ll pick up the keyboard or something. But all I get is a bounding box around my head. And if I move around it still doesn’t pick up anything else besides my head.

This is how I called detectnet with this custom model

$ /usr/local/bin/detectnet --model=SSD-Mobilenet-v1-Head/ssd-mobilenet.onnx --labels=../SSD-Mobilenet-v1-Head/ssd_coco_labels.txt --input-blob=input_0 --output-cvg=scores --output-bbox=boxes

My questions are:

  1. Why am I only seeing this one class and bounding box drawn if the training was based on the model that had 80 or more classes? My labels.txt when I run detectnet has all of the classes including the new “head” class.
  2. Where can I get a copy of the SSD-Mobilenet-v2 model to apply transfer learning to? I like the performance of this model more than v1.

If it’s relevant, here are my exact steps that I used to setup my environment for training the model, since I used a container on my laptop instead of training on the jetson using your jetson-inference container:

Based on your recommendation at can-i-use-my-containers-from-jetson-nano-on-windows-or-linux I used a container I found at pytorch-containers to do the training on WSL2 in a container.
Inside of the container (docker run --gpus all -it nvcr.io/nvidia/pytorch:22.01-py3), I cloned your pytorch repo inside of the container from here dusty-nv/pytorch-ssd and then followed the instructions at pytorch-ssd-transfer-learning-instructions to download the ssd-mobilenet.pth and download the training set for “Human head”, and trained it with a batch of 4 and for 30 epochs. The training was successful, and then I converted the model to ONNX format and copied it onto my jetson nano and ran detectnet with it. I think the pytorch version was 1.1, but don’t think that should matter since the transfer learning process should give the same results.

Hi,

Could you share the steps to retrain the model?
If the detection for other classes is required, please also add them to the training process.

Thanks.

Hi @AastaLLL,

  1. First I download the model:
wget https://nvidia.box.com/shared/static/djf5w54rjvpqocsiztzaandq1m3avr7c.pth -O models/ssd-head/mobilenet-v1-ssd-mp-0_675.pth
  1. Then I download the images:
$ python3 open_images_downloader.py --max-images=10000 --class-names "Human head,Ball" --data=data/head
  1. Then I train the model:
python3 train_ssd.py --data=data/head --model-dir=models/ssd-head --batch-size=4 --epochs=30
  1. My “labels.txt” during training is
BACKGROUND
Human head
Ball

I was under the impression that this model is already pretrained on the PASCAL VOC data set, and that by training the model in this way I should have the ability to detect everything that model was trained on, plus the “Human head” and “Ball” class. Cause then when I actually use the newly trained detection model, my labels.txt has “Human head, ball, Person, cat, dog, etc”.

Hi,

Detection includes the feature extractor + fully connected layer to the final class.

The model trains with a large dataset, the feature extractor should still have the capability to detect various objects.
But the fully connected part is rewritten with three classes in the training process.

If you want to detect all the classes in the PASCAL dataset, please add them to the output class.

Thanks.

I’m not exactly sure what you mean. I’m not super familiar with some of the terminology you are using. My understanding is that the steps I took are called transfer learning, in which I’m adding new features for the model to detect by training the model on the new object data set only. And then when I execute the code with the new model I thought I just need to make sure all of the classes are added to the labels.txt file. So as far as I know, I should see all objects detected with this newly trained model, and I don’t understand why it’s not.

So I did add all of the classes to the output class, but the only thing being detected is the new class I trained on.

Hi,

Do you also download the database to train the original class as well?
Thanks.

No I haven’t downloaded the database to train the original classes. Everything I shared here are the exact steps I followed. I thought it was a pretrained model, and that I could follow these steps to ADD a few classes to it and keep the original detection capability in tact. Why would I need to have the original database if this is a pretrained model? I thought the point of pre-trained was that it’s already trained to detect certain classes, and then you just add onto that.

Hi,

It keeps the feature extractor but the classification part will be overwritten in the training process.

Thanks.

I realize that I was believing that transfer learning meant that you get to keep the classes it was trained on and add new ones. When in reality I just learned that transfer learning means that you train a model with a specific architecture for a specific problem type (classification of objects in images) and then rather than retraining all of the model parameters to classify different images, you keep the internal layers the same (optionally) and just train the output layers and end up with only those classes. I get it now. Thanks for your help!

This topic was automatically closed 14 days after the last reply. New replies are no longer allowed.