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

[datumbox]

Feature Improvement

The test_detection_model_trainable_backbone_layers test currently downloads the weights of the backbone:

vision/test/test_models.py

Line 786 in 6530546

pretrained=False, pretrained_backbone=True, trainable_backbone_layers=trainable_layers

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 was initially worried that the weight download would be a problem for fbcode since we can't have internet access, but it looks like it's working fine, I think the manifold weights are properly picked up (D31758310).

However these tests are very long and add up for an extra 8 minutes on the run (https://www.internalfb.com/intern/testinfra/testconsole/testrun/844425139291544/):

I wonder if it's the same on the CircleCI side? We might want to try to reduce the time

[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.

[NicolasHug]

Setting the value pretrained_backbone=True is necessary because the number of trainable layers depends on this value.

Dumb question: Could we just set the expected values of the test such that they don't depend on pretrained_backbone? Or would this make the test useless?

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.

Would it be enough to patch load_state_dict_from_url to return None and to also patch model.load_state_dict to be a no-op ?

[datumbox]

Could we just set the expected values of the test such that they don't depend on pretrained_backbone?

Unfortunately the expected values depend on pretrained_backbone. Basically if you pass False, then all weights should be trainable. The option of what to freeze makes sense only when we use pre-trained weights...

Would it be enough to patch load_state_dict_from_url to return None and to also patch model.load_state_dict to be a no-op ?

Yes. We can easily patch load_state_dict_from_url to return None because it's a single method of the Weights class. Unfortunately for model, it's not initialized yet to make the load_state_dict method no-op on the instance level but we will have to do it on the nn.Module class level. Do you think that's a problem? If you think that's viable, we can definitely give it a try.

Edit: To clarify what I meant with "the model isn't initialized yet".

model = resnet50(weights=ResNet50Weights.ImageNet1k_RefV1)

the loading of the weights happens within the method. So we can't yet overwrite the model.load_state_dict on the instance level.

[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]

I've tried:

mocked = [
    mocker.patch('torch.nn.Module.load_state_dict'),
    mocker.patch('torchvision._internally_replaced_utils.load_state_dict_from_url'),
    mocker.patch('torch.hub.load_state_dict_from_url'),
]
# test code goes here
assert all(m.call_count > 0 for m in mocked)

Though the load_state_dict is called, the other two aren't. I know that the weights are being downloaded. Any thoughts?

[NicolasHug]

I think this is because patch doesn't actually change the end object, just the "links" that point to it

patch() works by (temporarily) changing the object that a name points to with another one. There can be many names pointing to any individual object, so for patching to work you must ensure that you patch the name used by the system under test. The basic principle is that you patch where an object is looked up, which is not necessarily the same place as where it is defined

from https://docs.python.org/3/library/unittest.mock.html#where-to-patch

So we probably need to patch load_state_dict_from_url from the files where the actual backbones are loaded.

This seems to be enough:

    mocker.patch('torchvision.models.resnet.load_state_dict_from_url')
    mocker.patch('torchvision.models.mobilenetv3.load_state_dict_from_url')
    mocker.patch('torchvision.models.vgg.load_state_dict_from_url')
    mocker.patch('torchvision.models.detection.ssd.load_state_dict_from_url')
    mocker.patch('torch.nn.Module.load_state_dict')

but it's a bit depressing that we have to manually write all these.

[datumbox]

@NicolasHug Oh darn... This can easily be fixed on the new API where the entire loading is taken care by a single method, but the old one is problematic.

@pmeier Any idea if there is a better way to patch the calls on load_state_dict_from_url globally on the tests? We basically want to no-op both load_state_dict_from_url and load_state_dict.

[datumbox]

Example of flakiness caused by downloading the weights.