adding test for trainable paramters in detection models (#4632)

* adding test for trainable paramters in detection models

* modifying range of trainable layers

Co-authored-by: Vasilis Vryniotis <[email protected]>

Author: jdsgomes

test/test_models.py

@@ -267,6 +267,43 @@ def _check_input_backprop(model, inputs):
 }
 
 
+# The following contains configuration and expected values to be used tests that are model specific
+_model_tests_values = {
+    "retinanet_resnet50_fpn": {
+        "max_trainable": 5,
+        "n_trn_params_per_layer": [36, 46, 65, 78, 88, 89],
+    },
+    "keypointrcnn_resnet50_fpn": {
+        "max_trainable": 5,
+        "n_trn_params_per_layer": [48, 58, 77, 90, 100, 101],
+    },
+    "fasterrcnn_resnet50_fpn": {
+        "max_trainable": 5,
+        "n_trn_params_per_layer": [30, 40, 59, 72, 82, 83],
+    },
+    "maskrcnn_resnet50_fpn": {
+        "max_trainable": 5,
+        "n_trn_params_per_layer": [42, 52, 71, 84, 94, 95],
+    },
+    "fasterrcnn_mobilenet_v3_large_fpn": {
+        "max_trainable": 6,
+        "n_trn_params_per_layer": [22, 23, 44, 70, 91, 97, 100],
+    },
+    "fasterrcnn_mobilenet_v3_large_320_fpn": {
+        "max_trainable": 6,
+        "n_trn_params_per_layer": [22, 23, 44, 70, 91, 97, 100],
+    },
+    "ssd300_vgg16": {
+        "max_trainable": 5,
+        "n_trn_params_per_layer": [45, 51, 57, 63, 67, 71],
+    },
+    "ssdlite320_mobilenet_v3_large": {
+        "max_trainable": 6,
+        "n_trn_params_per_layer": [96, 99, 138, 200, 239, 257, 266],
+    },
+}
+
+
 def _make_sliced_model(model, stop_layer):
     layers = OrderedDict()
     for name, layer in model.named_children():
@@ -740,5 +777,18 @@ def test_quantized_classification_model(model_name):
         raise AssertionError(f"model cannot be scripted. Traceback = {str(tb)}") from e
 
 
+@pytest.mark.parametrize("model_name", get_available_detection_models())
+def test_detection_model_trainable_backbone_layers(model_name):
+    max_trainable = _model_tests_values[model_name]["max_trainable"]
+    n_trainable_params = []
+    for trainable_layers in range(0, max_trainable + 1):
+        model = torchvision.models.detection.__dict__[model_name](
+            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]))
+    assert n_trainable_params == _model_tests_values[model_name]["n_trn_params_per_layer"]
+
+
 if __name__ == "__main__":
     pytest.main([__file__])