Delete layers in Faster-RCNN

username: Jiajun_Chen
cooked: <p>How can I delete the fc7 layer in Faster RCNN?<br>
And the following is part of the model structure</p>
<pre><code class="lang-auto">(roi_heads): RoIHeads(
    (box_roi_pool): MultiScaleRoIAlign()
    (box_head): TwoMLPHead(
      (fc6): Linear(in_features=12544, out_features=1024, bias=True)
      (fc7): Linear(in_features=1024, out_features=1024, bias=True)
    )
    (box_predictor): FastRCNNPredictor(
      (cls_score): Linear(in_features=1024, out_features=91, bias=True)
      (bbox_pred): Linear(in_features=1024, out_features=364, bias=True)
    )
  )
</code></pre>
<p>And if I do:</p>
<pre><code class="lang-auto">model.roi_heads.box_head=model.roi_heads.box_head.fc6
</code></pre>
<p>The model size will becoe super huge. When I put model in cuda, the CUDA memory will go from 7GB to 15GB  immediately.<br>
What’s the issue here?</p>

username: ptrblck
cooked: <p>Are you seeing the increase in memory, if you push another model to the device or do you delete the first one?<br>
You approach looks valid. Alternatively you could try to replace the <code>fc7</code> layer with an <code>nn.Identity()</code> module.</p>

username: Jiajun_Chen
cooked: <p>Thanks, I will try that…<br>
For the code, I simply change the model first and then push it to cuda device. Compared to the original model, the modified version cause cuda memory went from 7GB to 15GB.<br>
And also if I add a dropout layer after fc6 layer by using model.add_module(Dropout), is completely fine</p>

username: ptrblck
cooked: <p>I haven’t checked the shapes, but I assume the memory might blow up, since the flatten operation is missing, which is defined in <code>TwoMLPHead</code> in <a href="https://github.com/pytorch/vision/blob/f9ef235c402f48a335293c626e17bd8504d3af87/torchvision/models/detection/faster_rcnn.py#L252">this line of code</a>.</p>
<p>Replace it with:</p>
<pre><code class="lang-python">model.roi_heads.box_head = nn.Sequential(nn.Flatten(start_dim=1), model.roi_heads.box_head.fc6)
</code></pre>
<p>and it should work fine.</p>

username: Jiajun_Chen
cooked: <p>Wow, Thanks!. It did work.<br>
I forget the flatten operation at first, what a shame</p>

username: monster
cooked: <p>How can I add a new layer fc8 instead of deleting last layer ?</p>
<blockquote>
<p>(roi_heads): RoIHeads(<br>
(box_roi_pool): MultiScaleRoIAlign()<br>
(box_head): TwoMLPHead(<br>
(fc6): Linear(in_features=12544, out_features=1024, bias=True)<br>
(fc7): Linear(in_features=1024, out_features=1024, bias=True)<br>
(fc8): Linear(in_features=1024, out_features=1024, bias=True)<br>
)<br>
(box_predictor): FastRCNNPredictor(<br>
(cls_score): Linear(in_features=1024, out_features=91, bias=True)<br>
(bbox_pred): Linear(in_features=1024, out_features=364, bias=True)<br>
)<br>
)</p>
</blockquote>

username: ptrblck
cooked: <p>You can replace the attribute directly. I’m not sure about the nesting of the posted code, but it could look like:</p>
<pre><code class="lang-python">model.roi_heads.box_head.fc8 = nn.Linear(1024, ...)
</code></pre>

username: NumesSanguis
cooked: <p>I don’t know if in earlier versions of PyTorch the following works, but in v1.6 deleting a layer is as simple as:</p>
<pre><code class="lang-python"># top level layer
del model.fc
# untested: nested layer
del model.roi_heads.box_head.fc8
</code></pre>
<p>This both removes the layer from model.modules and model.state_dict.<br>
This is also does not create zombie layers, as an Identity layer would do. Simplifying model loading.</p>

username: ptrblck
cooked: <p>This approach could also break the model, since the <code>forward</code> pass would still try to call the deleted module:</p>
<pre><code class="lang-python">class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc1 = nn.Linear(1, 1)
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Linear(1, 1)
        
    def forward(self, x):
        x = self.relu1(self.fc1(x))
        x = self.fc2(x)
        return x

model = MyModel()
x = torch.randn(1, 1)
out = model(x)

del model.fc2

out = model(x)
&gt; ModuleAttributeError: 'MyModel' object has no attribute 'fc2'
</code></pre>
<p>While this approach might work well to “clean” the <code>state_dict</code>, I would not recommend it if the model should be used further without rewriting the <code>forward</code> method.</p>