6
6
import torch
7
7
from common_utils import set_rng_seed
8
8
from torchvision import models
9
+ from torchvision .models import efficientnet , mobilenet , resnet
9
10
from torchvision .models ._utils import IntermediateLayerGetter
10
- from torchvision .models .detection .backbone_utils import mobilenet_backbone , resnet_fpn_backbone
11
+ from torchvision .models .detection .backbone_utils import (
12
+ BackboneWithFPN ,
13
+ efficientnet_fpn_backbone ,
14
+ mobilenet_backbone ,
15
+ resnet_fpn_backbone ,
16
+ )
11
17
from torchvision .models .feature_extraction import create_feature_extractor , get_graph_node_names
12
18
13
19
@@ -16,7 +22,7 @@ def get_available_models():
16
22
return [k for k , v in models .__dict__ .items () if callable (v ) and k [0 ].lower () == k [0 ] and k [0 ] != "_" ]
17
23
18
24
19
- @pytest .mark .parametrize ("backbone_name" , ( "resnet18" , "resnet50" ) )
25
+ @pytest .mark .parametrize ("backbone_name" , resnet . __all__ [ 1 :] )
20
26
def test_resnet_fpn_backbone (backbone_name ):
21
27
x = torch .rand (1 , 3 , 300 , 300 , dtype = torch .float32 , device = "cpu" )
22
28
y = resnet_fpn_backbone (backbone_name = backbone_name , pretrained = False )(x )
@@ -28,16 +34,34 @@ def test_resnet_fpn_backbone(backbone_name):
28
34
resnet_fpn_backbone (backbone_name , False , returned_layers = [0 , 1 , 2 , 3 ])
29
35
with pytest .raises (ValueError , match = r"Each returned layer should be in the range" ):
30
36
resnet_fpn_backbone (backbone_name , False , returned_layers = [2 , 3 , 4 , 5 ])
37
+ model = resnet_fpn_backbone (backbone_name , False )
38
+ assert isinstance (model , BackboneWithFPN )
31
39
32
40
33
- @pytest .mark .parametrize ("backbone_name" , ( "mobilenet_v2" , "mobilenet_v3_large" , "mobilenet_v3_small" ) )
41
+ @pytest .mark .parametrize ("backbone_name" , mobilenet . mv2_all [ 1 :] + mobilenet . mv3_all [ 1 :] )
34
42
def test_mobilenet_backbone (backbone_name ):
35
43
with pytest .raises (ValueError , match = r"Trainable layers should be in the range" ):
36
44
mobilenet_backbone (backbone_name = backbone_name , pretrained = False , fpn = False , trainable_layers = - 1 )
37
45
with pytest .raises (ValueError , match = r"Each returned layer should be in the range" ):
38
46
mobilenet_backbone (backbone_name , False , fpn = True , returned_layers = [- 1 , 0 , 1 , 2 ])
39
47
with pytest .raises (ValueError , match = r"Each returned layer should be in the range" ):
40
48
mobilenet_backbone (backbone_name , False , fpn = True , returned_layers = [3 , 4 , 5 , 6 ])
49
+ model_fpn = mobilenet_backbone (backbone_name , False , fpn = True )
50
+ assert isinstance (model_fpn , BackboneWithFPN )
51
+ model = mobilenet_backbone (backbone_name , False , fpn = False )
52
+ assert isinstance (model , torch .nn .Sequential )
53
+
54
+
55
+ @pytest .mark .parametrize ("backbone_name" , efficientnet .__all__ [1 :])
56
+ def test_efficientnet_fpn_backbone (backbone_name ):
57
+ with pytest .raises (ValueError , match = r"Trainable layers should be in the range" ):
58
+ efficientnet_fpn_backbone (backbone_name = backbone_name , pretrained = False , trainable_layers = - 1 )
59
+ with pytest .raises (ValueError , match = r"Each returned layer should be in the range" ):
60
+ efficientnet_fpn_backbone (backbone_name , False , returned_layers = [- 1 , 0 , 1 , 2 ])
61
+ with pytest .raises (ValueError , match = r"Each returned layer should be in the range" ):
62
+ efficientnet_fpn_backbone (backbone_name , False , returned_layers = [3 , 4 , 5 , 6 , 9 ])
63
+ model = efficientnet_fpn_backbone (backbone_name , False )
64
+ assert isinstance (model , BackboneWithFPN )
41
65
42
66
43
67
# Needed by TestFxFeatureExtraction.test_leaf_module_and_function
0 commit comments