Reduce the number of classes in Faster R-CNN

username: Luppo
cooked: <p>Hello everyone.</p>
<p>I am using Faster R-CNN for object detection.<br>
Since I only need it to detect vehicles, I am just filtering out labels of non-vehicle objects, however I would like the network to output scores and bounding boxes for vehicles only.</p>
<p>Mainly, I need to change the number of output features of <code>model.roi_heads.box_predictor.bbox_pred</code>. Currently, it is a linear layer with <code>out_features=364</code>, basically 4 outputs for each of the 91 classes of COCO dataset. However I would like to exploit the feature extraction at this stage to predict other values.</p>
<p>I’ll try to explain better with an example:</p>
<ol>
<li>Let’s say I am just interested in predicting ‘car’ class</li>
</ol>
<pre><code class="lang-auto">car_label = 3
</code></pre>
<ol start="2">
<li>I extract the 4 rows in the weights matrix to predict the bounding box (since I want to use pre-trained Faster R-CNN)</li>
</ol>
<pre><code class="lang-auto">my_weights = model.roi_heads.box_predictor.bbox_pred.weight.data[4*car_label:4*(car_label+1), :]
my_bias = model.roi_heads.box_predictor.bbox_pred.bias.data[4*car_label:4*(car_label+1)]
</code></pre>
<ol start="3">
<li>I want the bounding box predictor to regress another output (5 outputs instead of 4)</li>
</ol>
<pre><code class="lang-auto">model.roi_heads.box_predictor.bbox_pred= nn.Linear(in_features=1024, out_features=5, bias=True)
</code></pre>
<ol start="4">
<li>But I still want to have the bounding box, exploiting the pre-trained network</li>
</ol>
<pre><code class="lang-auto">model.roi_heads.box_predictor.bbox_pred.weight.data[:4,:] = my_weights
model.roi_heads.box_predictor.bbox_pred.bias.data[:4] = my_bias
</code></pre>
<p>This obviously doesn’t work:</p>
<pre><code class="lang-auto">RuntimeError: The expanded size of the tensor (1) must match the existing size (2) at non-singleton dimension 1.  Target sizes: [10000, 1].  Tensor sizes: [10000, 2]
</code></pre>
<p>But even if I set 4 output features only (<code>out_features=4</code>) an error still exists because it removes any lower score box from the output:</p>
<pre><code class="lang-auto">/usr/local/lib/python3.6/dist-packages/torchvision/models/detection/roi_heads.py in postprocess_detections(self, class_logits, box_regression, proposals, image_shapes)
    502             # remove low scoring boxes
    503             inds = torch.nonzero(scores &gt; self.score_thresh).squeeze(1)
--&gt; 504             boxes, scores, labels = boxes[inds], scores[inds], labels[inds]
    505 
    506             # remove empty boxes

IndexError: index is out of bounds for dimension with size 0
</code></pre>
<p>Any idea on how to do something like this or how to delete the removal of low scoring boxes (since I actually only need one box only)? I also accept ideas on how to add another linear layer having, as input, the 1024 <code>in_features</code> and providing a readable output, even though I prefer to delete this last large matrix multiplication since I don’t need it (and I tried to delete it following <a href="https://discuss.pytorch.org/t/how-to-delete-layer-in-pretrained-model/17648">these steps</a> but I had errors on the input of a <code>conv</code> layer).</p>
<p>Thank you.</p>