Fixing bug on SSD backbone freezing

[datumbox]

Fixes #4548

There was a bug on the way we counted the frozen layers. This mean that the first "block" (layers 0-3 inclusive) of VGG was frozen. Here is how to confirm this is the case on the current main branch:

from torchvision.models.detection import *

freeze_before = 4
trained = ssd300_vgg16(pretrained=True)
nottrained = ssd300_vgg16(pretrained=False, pretrained_backbone=True)
sum_trained = sum(x.sum() for x in trained.backbone.features[:freeze_before].parameters())
sum_nottrained = sum(x.sum() for x in nottrained.backbone.features[:freeze_before].parameters())

assert sum_trained == sum_nottrained

This is how we can confirm that we freeze the right number of params:

import torchvision

n_trainable_params = []
for trainable_layers in range(6):
    model = torchvision.models.detection.ssd300_vgg16(pretrained=False,
            pretrained_backbone=True,
            trainable_backbone_layers=trainable_layers)

    n_trainable_params.append(len([p for p in model.parameters() if p.requires_grad]))

print("SSD:")
print(f"Number of trainable parameters: {n_trainable_params}")
print(f"Total number of parameters: {len([p for p in model.parameters()])}")

Results:

SSD:
Number of trainable parameters: [45, 51, 57, 63, 67, 71]
Total number of parameters: 71

cc @datumbox

[NicolasHug]

Stamping

[NicolasHug]

As a follow up, should we consider including your snippets above as tests?

[datumbox]

@NicolasHug Thanks for the stamps.

I definitely agree we should write tests for the trainable_backbone_layers parameter of all detection models. The challenge is that re-initializing the models is quite slow and likely to impact significantly our execution speed. I'll create a separate issue to investigate how we could do this.

Edit: I've opened an issue at #4592