How reduce number of classes in faster rcnn

username: ayni-af
cooked: <p>Hi,</p>
<p>I am new in the field of object detection, I will be grateful if you could help me to reduce the number of detected objects in a pre-trained model that is trained on the coco dataset. I want only to detect “person” and “dog”.</p>
<p>I am using fasterrcnn_resnet50_fpn model:</p>
<blockquote>
<p><span class="hashtag-raw">#load</span> mode<br>
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)</p>
</blockquote>
<p>I am not sure who should I modify the model to only detect class 1 (person) and class 18 (dog). I do not want to train the model on new data.</p>
<p>I will appreciate it if you can help me with this problem.<br>
Thank you</p>

username: Dwight_Foster
cooked: <p>In the call you can just set num_classes to whatever you want. If you go in the docs <a href="https://pytorch.org/vision/stable/models.html#object-detection-instance-segmentation-and-person-keypoint-detection" rel="noopener nofollow ugc">here</a> it shows you how to do this as well as other things.</p>

username: JamesDickens
cooked: <p>Keep in mind that if you want to use the coco-pretrained backbone and rpn, and then train with a new ROI head (for different class structure), you can simply take the backbone and rpn, and use those to initialize a new faster r-cnn module by passing in the backbone and rpn as arguments to faster-rcnn in torchvision.</p>

username: ayni-af
cooked: <p>Hi <a class="mention" href="/u/dwight_foster">@Dwight_Foster</a> ,</p>
<p>thank you for your reply. But how can specify that I am only interested in class 1 (person) &amp; class 18 (dog)? I do not want to train again the fasterrcnn_resnet50_fpn model. I want that the model only detects my wished classes and ingnores other 88 classes.</p>
<p>I tried to set the num_classes to 3 but get an error.<br>
My code:</p>
<pre><code class="lang-auto">coco_names = ['__background__', 'person', 'dog']

model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True, num_classes=len(coco_names))


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.eval().to(device)

path='Downloads/input2/test.jpg'
image = Image.open(path)

boxes, classes, labels = detect_utils.predict(image, model, device, 0.8)
image =detect_utils.draw_boxes(boxes, classes, labels, image)
#cv2.imshow('Image', image)
save_name = f"{path.split('/')[-1].split('.')[0]}"
cv2.imwrite(f"outputs/{save_name}.jpg", image)
</code></pre>
<p>the error that I get:</p>
<pre><code class="lang-auto">RuntimeError: Error(s) in loading state_dict for FasterRCNN:
	size mismatch for roi_heads.box_predictor.cls_score.weight: copying a param with shape torch.Size([91, 1024]) from checkpoint, the shape in current model is torch.Size([3, 1024]).
	size mismatch for roi_heads.box_predictor.cls_score.bias: copying a param with shape torch.Size([91]) from checkpoint, the shape in current model is torch.Size([3]).
	size mismatch for roi_heads.box_predictor.bbox_pred.weight: copying a param with shape torch.Size([364, 1024]) from checkpoint, the shape in current model is torch.Size([12, 1024]).
	size mismatch for roi_heads.box_predictor.bbox_pred.bias: copying a param with shape torch.Size([364]) from checkpoint, the shape in current model is torch.Size([12]).
</code></pre>

username: ayni-af
cooked: <p>Hi <a class="mention" href="/u/jamesdickens">@JamesDickens</a> ,</p>
<p>Thank you for your answer. I do not want to train the model with a new dataset. I want just that model detects my wishes classes which are available in the coco dataset:</p>
<pre><code class="lang-auto">coco_names = ['__background__', 'person', 'dog']
</code></pre>
<p>and ignores other classes.<br>
I am tried to set the num_classes to 3 by loading the model it did not work.</p>

username: Dwight_Foster
cooked: <p>Ok the error is because you cannot load in the pretrained model and set num classes to a different value because then the weights won’t match. Or at least it didn’t work for me either. You can try this however:</p>
<pre><code class="lang-auto">model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
model.roi_heads.box_predictor.cls_score = nn.Linear(1024,len(coco_names)
</code></pre>
<p>that should work.</p>

username: duddal
cooked: <p><a class="mention" href="/u/dwight_foster">@Dwight_Foster</a> Hi, I know it’s been some time since this post has been active.<br>
But I tried your method and I have some doubts:</p>
<ol>
<li>
<code>model.roi_heads.box_predictor.cls_score = nn.Linear(1024,len(coco_names)</code>. Here we are just telling our model to predict for 3 classes but how does the model know that the classes should be ‘background’, ‘person’ and a ‘dog’?</li>
</ol>

username: Dwight_Foster
cooked: <p>When you initialize the model it does not know which class is which. That is the whole point of training the model is trying to learn which classes are which.</p>