Improve test of backbone utils

Author: talregev

test/test_backbone_utils.py

@@ -7,7 +7,7 @@
 from common_utils import set_rng_seed
 from torchvision import models
 from torchvision.models._utils import IntermediateLayerGetter
-from torchvision.models.detection.backbone_utils import mobilenet_backbone, resnet_fpn_backbone
+from torchvision.models.detection.backbone_utils import BackboneWithFPN, mobilenet_backbone, resnet_fpn_backbone
 from torchvision.models.feature_extraction import create_feature_extractor, get_graph_node_names
 
 
@@ -16,10 +16,12 @@ def get_available_models():
     return [k for k, v in models.__dict__.items() if callable(v) and k[0].lower() == k[0] and k[0] != "_"]
 
 
-@pytest.mark.parametrize("backbone_name", ("resnet18", "resnet50"))
+@pytest.mark.parametrize("backbone_name", ("resnet18", "resnet50", "resnet101"))
 def test_resnet_fpn_backbone(backbone_name):
     x = torch.rand(1, 3, 300, 300, dtype=torch.float32, device="cpu")
-    y = resnet_fpn_backbone(backbone_name=backbone_name, pretrained=False)(x)
+    model = resnet_fpn_backbone(backbone_name=backbone_name, pretrained=False)
+    assert isinstance(model, BackboneWithFPN)
+    y = model(x)
     assert list(y.keys()) == ["0", "1", "2", "3", "pool"]
 
     with pytest.raises(ValueError, match=r"Trainable layers should be in the range"):
@@ -38,6 +40,10 @@ def test_mobilenet_backbone(backbone_name):
         mobilenet_backbone(backbone_name, False, fpn=True, returned_layers=[-1, 0, 1, 2])
     with pytest.raises(ValueError, match=r"Each returned layer should be in the range"):
         mobilenet_backbone(backbone_name, False, fpn=True, returned_layers=[3, 4, 5, 6])
+    model_fpn = mobilenet_backbone(backbone_name, False, fpn=True)
+    assert isinstance(model_fpn, BackboneWithFPN)
+    model = mobilenet_backbone(backbone_name, False, fpn=False)
+    assert isinstance(model, torch.nn.Sequential)
 
 
 # Needed by TestFxFeatureExtraction.test_leaf_module_and_function