`mobilenet_backbone` behavior when `trainable_layers = 0`

[ekorman-strive]

🐛 Bug

When torchvision.models.detection.backbone_utils.mobilenet_backbone is called with trainable_layers=0, more parameters
are trainable than I would have expected.

To Reproduce

The following script:

from torchvision.models.detection.backbone_utils import (
    mobilenet_backbone,
    resnet_fpn_backbone,
)

n_trainable_params = []
for trainable_layers in range(6):
    net = mobilenet_backbone(
        "mobilenet_v3_large",
        pretrained=False,
        fpn=True,
        trainable_layers=trainable_layers,
    )

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

print(f"Mobilenet number of trainable parameters: {n_trainable_params}")

n_trainable_params = []
for trainable_layers in range(6):
    net = resnet_fpn_backbone(
        "resnet50", pretrained=False, trainable_layers=trainable_layers,
    )

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

print(f"Resnet number of trainable parameters: {n_trainable_params}")

outputs

Mobilenet number of trainable parameters: [63, 9, 30, 56, 77, 83]
Resnet number of trainable parameters: [16, 26, 45, 58, 68, 69]

However I would expect the first list to be increasing (like the second one). This line

vision/torchvision/models/detection/backbone_utils.py

Line 154 in 9c31d1d

freeze_before = num_stages if trainable_layers == 0 else stage_indices[num_stages - trainable_layers]

looks wrong since num_stages isn't correlated to an index of backbone.

Expected behavior

The number of trainable parameters for a mobilenet backbone when trainable_layers is set to 0 should be less than when trainable_layers is set to anything greater than 0.

Environment

PyTorch version: 1.8.1+cu102
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.5 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.10

Python version: 3.7 (64-bit runtime)
Python platform: Linux-5.4.72-microsoft-standard-WSL2-x86_64-with-debian-buster-sid
Is CUDA available: False
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.17.4
[pip3] torch==1.8.1
[pip3] torchvision==0.9.1
[conda] numpy                     1.17.4                   pypi_0    pypi
[conda] torch                     1.8.1                    pypi_0    pypi
[conda] torchvision               0.9.1                    pypi_0    pypi

[datumbox]

@ekorman-strive Thanks for reporting, you are absolutely right. This is a bug. I'll follow up with a PR shortly.