Delete layers in Faster-RCNN

How can I delete the fc7 layer in Faster RCNN?
And the following is part of the model structure

(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)
    )
  )

And if I do:

model.roi_heads.box_head=model.roi_heads.box_head.fc6

The model size will becoe super huge. When I put model in cuda, the CUDA memory will go from 7GB to 15GB immediately.
What’s the issue here?

Are you seeing the increase in memory, if you push another model to the device or do you delete the first one?
You approach looks valid. Alternatively you could try to replace the fc7 layer with an nn.Identity() module.

Thanks, I will try that…
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.
And also if I add a dropout layer after fc6 layer by using model.add_module(Dropout), is completely fine

I haven’t checked the shapes, but I assume the memory might blow up, since the flatten operation is missing, which is defined in TwoMLPHead in this line of code.

Replace it with:

model.roi_heads.box_head = nn.Sequential(nn.Flatten(start_dim=1), model.roi_heads.box_head.fc6)

and it should work fine.

2 Likes

Wow, Thanks!. It did work.
I forget the flatten operation at first, what a shame

1 Like

How can I add a new layer fc8 instead of deleting last layer ?

(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)
(fc8): 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)
)
)

You can replace the attribute directly. I’m not sure about the nesting of the posted code, but it could look like:

model.roi_heads.box_head.fc8 = nn.Linear(1024, ...)
1 Like

I don’t know if in earlier versions of PyTorch the following works, but in v1.6 deleting a layer is as simple as:

# top level layer
del model.fc
# untested: nested layer
del model.roi_heads.box_head.fc8

This both removes the layer from model.modules and model.state_dict.
This is also does not create zombie layers, as an Identity layer would do. Simplifying model loading.

This approach could also break the model, since the forward pass would still try to call the deleted module:

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)
> ModuleAttributeError: 'MyModel' object has no attribute 'fc2'

While this approach might work well to “clean” the state_dict, I would not recommend it if the model should be used further without rewriting the forward method.