The `test_detection_model_trainable_backbone_layers` test shouldn't download the pretrained_backbone weights

Feature Improvement

The test_detection_model_trainable_backbone_layers test currently downloads the weights of the backbone: https://github.com/pytorch/vision/blob/6530546080148cd5d9533b50a8258f12cb050c82/test/test_models.py#L786

Setting the value pretrained_backbone=True is necessary because the number of trainable layers depends on this value. Unfortunately downloading pre-trained weights can lead to flakiness and slow tests and should be avoided. Until we setup a cache to store the weights locally on the CI, we should find a way to skip the actual downloading of weights during the test execution.

cc @datumbox @pmeier

[NicolasHug]

I think it should be possible to patch load_state_dict at the nn.Module level as you suggested, something like:

@pytest.mark.parametrize("model_name", get_available_detection_models())
def test_mock_method(model_name, mocker):

    mocker.patch('torch.nn.Module.load_state_dict')
    mocker.patch('torch.hub.load_state_dict_from_url')
 
    model = torchvision.models.detection.__dict__[model_name](
        pretrained=False, pretrained_backbone=True, trainable_backbone_layers=4,
    )

I haven’t checked, but this shouldn’t do any network call – this can be verified with pytest-sockets https://github.com/miketheman/pytest-socket

[datumbox]

@NicolasHug thanks for checking. Yes the weights are always added in the manifold so this won’t be a problem.

The speed is not as bad on CircleCI. We checked during merge (see here) and only 1 test appeared on the top slowest with execution time 14-15 secs. All others executed in less than 8 sec:

==== ========================= slowest 20 durations =============================
54.32s call     test/test_models.py::test_quantized_classification_model[mobilenet_v3_large]
39.12s call     test/test_models.py::test_quantized_classification_model[resnext101_32x8d]
38.79s call     test/test_models.py::test_quantized_classification_model[mobilenet_v2]
31.24s call     test/test_models.py::test_quantized_classification_model[shufflenet_v2_x0_5]
18.17s call     test/test_models.py::test_quantized_classification_model[googlenet]
15.27s call     test/test_models.py::test_quantized_classification_model[resnet50]
14.98s call     test/test_models.py::test_quantized_classification_model[shufflenet_v2_x1_0]
14.80s call     test/test_models.py::test_quantized_classification_model[shufflenet_v2_x2_0]
14.48s call     test/test_models.py::test_detection_model_trainable_backbone_layers[ssd300_vgg16]
14.10s call     test/test_models.py::test_quantized_classification_model[shufflenet_v2_x1_5]
13.88s call     test/test_datasets.py::LFWPairsTestCase::test_transforms
13.29s call     test/test_models.py::test_detection_model[cpu-fasterrcnn_mobilenet_v3_large_fpn]
12.53s call     test/test_models.py::test_classification_model[cpu-densenet201]
12.14s call     test/test_models.py::test_classification_model[cpu-densenet161]
11.28s call     test/test_models.py::test_classification_model[cpu-efficientnet_b7]
10.74s call     test/test_models.py::test_classification_model[cpu-densenet169]
10.66s call     test/test_backbone_utils.py::TestFxFeatureExtraction::test_jit_forward_backward[efficientnet_b7]
10.14s call     test/test_models.py::test_classification_model[cpu-regnet_y_32gf]
9.10s call     test/test_models.py::test_classification_model[cpu-efficientnet_b6]
8.80s call     test/test_datasets.py::LFWPeopleTestCase::test_transforms
=========================== short test summary info ============================

The problem can be solved quite easily using the new weights API, because you can easily patch the model loading method of the weights during tests and avoid their actual downloading. That’s a bit harder using the old pretrained approach. Any thoughts/ideas about this? I’m open to reverting if a good solution doesn’t exist now and we could add the test back once we’ve moved the multi-pretrained model mechanism in main.

EDIT:

I checked the new proposed API and patching it won’t be as easy as I remembered. I think we need to make some additional minor adjustments. Here is how the weights are typically loaded based on the current proposal:

if weights is not None:
    model.load_state_dict(weights.state_dict(progress=progress))

One could easily patch the state_dict method of the weights in the tests, so that they don’t actually download anything. Unfortunately, passing None on load_state_dict() wont work. We could have an extra step that checks that the weights are not None before loading them but this might require some extra thought.