test_extended_models.py 15.8 KB
Newer Older
1
import copy
Philip Meier's avatar
Philip Meier committed
2
import os
3
import pickle
4
5

import pytest
6
import test_models as TM
7
import torch
Nicolas Hug's avatar
Nicolas Hug committed
8
from common_extended_utils import get_file_size_mb, get_ops
9
from torchvision import models
10
from torchvision.models import get_model_weights, Weights, WeightsEnum
11
from torchvision.models._utils import handle_legacy_interface
12
from torchvision.models.detection.backbone_utils import mobilenet_backbone, resnet_fpn_backbone
13

14
15
16
17
run_if_test_with_extended = pytest.mark.skipif(
    os.getenv("PYTORCH_TEST_WITH_EXTENDED", "0") != "1",
    reason="Extended tests are disabled by default. Set PYTORCH_TEST_WITH_EXTENDED=1 to run them.",
)
18
19


20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
@pytest.mark.parametrize(
    "name, model_class",
    [
        ("resnet50", models.ResNet),
        ("retinanet_resnet50_fpn_v2", models.detection.RetinaNet),
        ("raft_large", models.optical_flow.RAFT),
        ("quantized_resnet50", models.quantization.QuantizableResNet),
        ("lraspp_mobilenet_v3_large", models.segmentation.LRASPP),
        ("mvit_v1_b", models.video.MViT),
    ],
)
def test_get_model(name, model_class):
    assert isinstance(models.get_model(name), model_class)


35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
@pytest.mark.parametrize(
    "name, model_fn",
    [
        ("resnet50", models.resnet50),
        ("retinanet_resnet50_fpn_v2", models.detection.retinanet_resnet50_fpn_v2),
        ("raft_large", models.optical_flow.raft_large),
        ("quantized_resnet50", models.quantization.resnet50),
        ("lraspp_mobilenet_v3_large", models.segmentation.lraspp_mobilenet_v3_large),
        ("mvit_v1_b", models.video.mvit_v1_b),
    ],
)
def test_get_model_builder(name, model_fn):
    assert models.get_model_builder(name) == model_fn


50
51
52
53
54
55
56
57
58
59
60
61
62
@pytest.mark.parametrize(
    "name, weight",
    [
        ("resnet50", models.ResNet50_Weights),
        ("retinanet_resnet50_fpn_v2", models.detection.RetinaNet_ResNet50_FPN_V2_Weights),
        ("raft_large", models.optical_flow.Raft_Large_Weights),
        ("quantized_resnet50", models.quantization.ResNet50_QuantizedWeights),
        ("lraspp_mobilenet_v3_large", models.segmentation.LRASPP_MobileNet_V3_Large_Weights),
        ("mvit_v1_b", models.video.MViT_V1_B_Weights),
    ],
)
def test_get_model_weights(name, weight):
    assert models.get_model_weights(name) == weight
63
64


65
66
67
68
69
70
71
72
73
74
75
76
77
@pytest.mark.parametrize("copy_fn", [copy.copy, copy.deepcopy])
@pytest.mark.parametrize(
    "name",
    [
        "resnet50",
        "retinanet_resnet50_fpn_v2",
        "raft_large",
        "quantized_resnet50",
        "lraspp_mobilenet_v3_large",
        "mvit_v1_b",
    ],
)
def test_weights_copyable(copy_fn, name):
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
    for weights in list(models.get_model_weights(name)):
        # It is somewhat surprising that (deep-)copying is an identity operation here, but this is the default behavior
        # of enums: https://docs.python.org/3/howto/enum.html#enum-members-aka-instances
        # Checking for equality, i.e. `==`, is sufficient (and even preferable) for our use case, should we need to drop
        # support for the identity operation in the future.
        assert copy_fn(weights) is weights


@pytest.mark.parametrize(
    "name",
    [
        "resnet50",
        "retinanet_resnet50_fpn_v2",
        "raft_large",
        "quantized_resnet50",
        "lraspp_mobilenet_v3_large",
        "mvit_v1_b",
    ],
)
def test_weights_deserializable(name):
    for weights in list(models.get_model_weights(name)):
        # It is somewhat surprising that deserialization is an identity operation here, but this is the default behavior
        # of enums: https://docs.python.org/3/howto/enum.html#enum-members-aka-instances
        # Checking for equality, i.e. `==`, is sufficient (and even preferable) for our use case, should we need to drop
        # support for the identity operation in the future.
        assert pickle.loads(pickle.dumps(weights)) is weights
104
105


106
107
108
109
110
111
112
@pytest.mark.parametrize(
    "module", [models, models.detection, models.quantization, models.segmentation, models.video, models.optical_flow]
)
def test_list_models(module):
    def get_models_from_module(module):
        return [
            v.__name__
113
            for k, v in module.__dict__.items()
114
115
116
117
118
119
120
121
            if callable(v) and k[0].islower() and k[0] != "_" and k not in models._api.__all__
        ]

    a = set(get_models_from_module(module))
    b = set(x.replace("quantized_", "") for x in models.list_models(module))

    assert len(b) > 0
    assert a == b
122
123


124
@pytest.mark.parametrize(
125
    "name, weight",
126
    [
127
128
        ("ResNet50_Weights.IMAGENET1K_V1", models.ResNet50_Weights.IMAGENET1K_V1),
        ("ResNet50_Weights.DEFAULT", models.ResNet50_Weights.IMAGENET1K_V2),
129
        (
130
131
            "ResNet50_QuantizedWeights.DEFAULT",
            models.quantization.ResNet50_QuantizedWeights.IMAGENET1K_FBGEMM_V2,
132
        ),
133
        (
134
135
            "ResNet50_QuantizedWeights.IMAGENET1K_FBGEMM_V1",
            models.quantization.ResNet50_QuantizedWeights.IMAGENET1K_FBGEMM_V1,
136
        ),
137
138
    ],
)
139
140
def test_get_weight(name, weight):
    assert models.get_weight(name) == weight
141
142


143
144
@pytest.mark.parametrize(
    "model_fn",
145
146
147
148
149
150
    TM.list_model_fns(models)
    + TM.list_model_fns(models.detection)
    + TM.list_model_fns(models.quantization)
    + TM.list_model_fns(models.segmentation)
    + TM.list_model_fns(models.video)
    + TM.list_model_fns(models.optical_flow),
151
152
)
def test_naming_conventions(model_fn):
153
    weights_enum = get_model_weights(model_fn)
154
    assert weights_enum is not None
155
    assert len(weights_enum) == 0 or hasattr(weights_enum, "DEFAULT")
156
157


158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
detection_models_input_dims = {
    "fasterrcnn_mobilenet_v3_large_320_fpn": (320, 320),
    "fasterrcnn_mobilenet_v3_large_fpn": (800, 800),
    "fasterrcnn_resnet50_fpn": (800, 800),
    "fasterrcnn_resnet50_fpn_v2": (800, 800),
    "fcos_resnet50_fpn": (800, 800),
    "keypointrcnn_resnet50_fpn": (1333, 1333),
    "maskrcnn_resnet50_fpn": (800, 800),
    "maskrcnn_resnet50_fpn_v2": (800, 800),
    "retinanet_resnet50_fpn": (800, 800),
    "retinanet_resnet50_fpn_v2": (800, 800),
    "ssd300_vgg16": (300, 300),
    "ssdlite320_mobilenet_v3_large": (320, 320),
}


174
175
@pytest.mark.parametrize(
    "model_fn",
176
177
178
179
180
181
    TM.list_model_fns(models)
    + TM.list_model_fns(models.detection)
    + TM.list_model_fns(models.quantization)
    + TM.list_model_fns(models.segmentation)
    + TM.list_model_fns(models.video)
    + TM.list_model_fns(models.optical_flow),
182
)
183
@run_if_test_with_extended
184
def test_schema_meta_validation(model_fn):
185
186
187
188
189
190
    # list of all possible supported high-level fields for weights meta-data
    permitted_fields = {
        "backend",
        "categories",
        "keypoint_names",
        "license",
191
        "_metrics",
192
        "min_size",
193
        "min_temporal_size",
194
195
196
        "num_params",
        "recipe",
        "unquantized",
197
        "_docs",
198
        "_ops",
Nicolas Hug's avatar
Nicolas Hug committed
199
        "_file_size",
200
201
    }
    # mandatory fields for each computer vision task
202
    classification_fields = {"categories", ("_metrics", "ImageNet-1K", "acc@1"), ("_metrics", "ImageNet-1K", "acc@5")}
203
    defaults = {
Nicolas Hug's avatar
Nicolas Hug committed
204
        "all": {"_metrics", "min_size", "num_params", "recipe", "_docs", "_file_size", "_ops"},
205
        "models": classification_fields,
206
        "detection": {"categories", ("_metrics", "COCO-val2017", "box_map")},
207
        "quantization": classification_fields | {"backend", "unquantized"},
208
209
210
211
212
213
        "segmentation": {
            "categories",
            ("_metrics", "COCO-val2017-VOC-labels", "miou"),
            ("_metrics", "COCO-val2017-VOC-labels", "pixel_acc"),
        },
        "video": {"categories", ("_metrics", "Kinetics-400", "acc@1"), ("_metrics", "Kinetics-400", "acc@5")},
214
        "optical_flow": set(),
215
    }
216
    model_name = model_fn.__name__
217
    module_name = model_fn.__module__.split(".")[-2]
218
    expected_fields = defaults["all"] | defaults[module_name]
219

220
    weights_enum = get_model_weights(model_fn)
221
222
    if len(weights_enum) == 0:
        pytest.skip(f"Model '{model_name}' doesn't have any pre-trained weights.")
223
224

    problematic_weights = {}
225
    incorrect_meta = []
226
    bad_names = []
227
    for w in weights_enum:
228
229
230
231
232
233
234
        actual_fields = set(w.meta.keys())
        actual_fields |= set(
            ("_metrics", dataset, metric_key)
            for dataset in w.meta.get("_metrics", {}).keys()
            for metric_key in w.meta.get("_metrics", {}).get(dataset, {}).keys()
        )
        missing_fields = expected_fields - actual_fields
235
236
237
        unsupported_fields = set(w.meta.keys()) - permitted_fields
        if missing_fields or unsupported_fields:
            problematic_weights[w] = {"missing": missing_fields, "unsupported": unsupported_fields}
238
239

        if w == weights_enum.DEFAULT or any(w.meta[k] != weights_enum.DEFAULT.meta[k] for k in ["num_params", "_ops"]):
240
            if module_name == "quantization":
241
                # parameters() count doesn't work well with quantization, so we check against the non-quantized
242
                unquantized_w = w.meta.get("unquantized")
243
244
245
246
247
248
249
250
251
                if unquantized_w is not None:
                    if w.meta.get("num_params") != unquantized_w.meta.get("num_params"):
                        incorrect_meta.append((w, "num_params"))

                    # the methodology for quantized ops count doesn't work as well, so we take unquantized FLOPs
                    # instead
                    if w.meta["_ops"] != unquantized_w.meta.get("_ops"):
                        incorrect_meta.append((w, "_ops"))

252
            else:
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
                # loading the model and using it for parameter and ops verification
                model = model_fn(weights=w)

                if w.meta.get("num_params") != sum(p.numel() for p in model.parameters()):
                    incorrect_meta.append((w, "num_params"))

                kwargs = {}
                if model_name in detection_models_input_dims:
                    # detection models have non default height and width
                    height, width = detection_models_input_dims[model_name]
                    kwargs = {"height": height, "width": width}

                calculated_ops = get_ops(model=model, weight=w, **kwargs)
                if calculated_ops != w.meta["_ops"]:
                    incorrect_meta.append((w, "_ops"))

269
270
        if not w.name.isupper():
            bad_names.append(w)
271

Nicolas Hug's avatar
Nicolas Hug committed
272
273
        if get_file_size_mb(w) != w.meta.get("_file_size"):
            incorrect_meta.append((w, "_file_size"))
274

275
    assert not problematic_weights
276
    assert not incorrect_meta
277
    assert not bad_names
278
279


280
@pytest.mark.parametrize(
281
    "model_fn",
282
283
284
285
286
287
    TM.list_model_fns(models)
    + TM.list_model_fns(models.detection)
    + TM.list_model_fns(models.quantization)
    + TM.list_model_fns(models.segmentation)
    + TM.list_model_fns(models.video)
    + TM.list_model_fns(models.optical_flow),
288
)
289
290
291
@run_if_test_with_extended
def test_transforms_jit(model_fn):
    model_name = model_fn.__name__
292
    weights_enum = get_model_weights(model_fn)
293
294
295
    if len(weights_enum) == 0:
        pytest.skip(f"Model '{model_name}' doesn't have any pre-trained weights.")

296
    defaults = {
297
298
299
        "models": {
            "input_shape": (1, 3, 224, 224),
        },
300
301
302
        "detection": {
            "input_shape": (3, 300, 300),
        },
303
304
305
        "quantization": {
            "input_shape": (1, 3, 224, 224),
        },
306
307
308
309
        "segmentation": {
            "input_shape": (1, 3, 520, 520),
        },
        "video": {
310
            "input_shape": (1, 3, 4, 112, 112),
311
        },
312
313
314
        "optical_flow": {
            "input_shape": (1, 3, 128, 128),
        },
315
    }
316
    module_name = model_fn.__module__.split(".")[-2]
317

318
319
320
    kwargs = {**defaults[module_name], **TM._model_params.get(model_name, {})}
    input_shape = kwargs.pop("input_shape")
    x = torch.rand(input_shape)
321
    if module_name == "optical_flow":
322
        args = (x, x)
323
    else:
324
325
        if module_name == "video":
            x = x.permute(0, 2, 1, 3, 4)
326
        args = (x,)
327

328
329
330
331
332
333
334
    problematic_weights = []
    for w in weights_enum:
        transforms = w.transforms()
        try:
            TM._check_jit_scriptable(transforms, args)
        except Exception:
            problematic_weights.append(w)
335

336
    assert not problematic_weights
Philip Meier's avatar
Philip Meier committed
337
338
339
340
341


# With this filter, every unexpected warning will be turned into an error
@pytest.mark.filterwarnings("error")
class TestHandleLegacyInterface:
342
    class ModelWeights(WeightsEnum):
Philip Meier's avatar
Philip Meier committed
343
344
345
346
347
348
349
        Sentinel = Weights(url="https://pytorch.org", transforms=lambda x: x, meta=dict())

    @pytest.mark.parametrize(
        "kwargs",
        [
            pytest.param(dict(), id="empty"),
            pytest.param(dict(weights=None), id="None"),
350
            pytest.param(dict(weights=ModelWeights.Sentinel), id="Weights"),
Philip Meier's avatar
Philip Meier committed
351
352
353
        ],
    )
    def test_no_warn(self, kwargs):
354
        @handle_legacy_interface(weights=("pretrained", self.ModelWeights.Sentinel))
Philip Meier's avatar
Philip Meier committed
355
356
357
358
359
360
361
        def builder(*, weights=None):
            pass

        builder(**kwargs)

    @pytest.mark.parametrize("pretrained", (True, False))
    def test_pretrained_pos(self, pretrained):
362
        @handle_legacy_interface(weights=("pretrained", self.ModelWeights.Sentinel))
Philip Meier's avatar
Philip Meier committed
363
364
365
366
367
368
369
370
        def builder(*, weights=None):
            pass

        with pytest.warns(UserWarning, match="positional"):
            builder(pretrained)

    @pytest.mark.parametrize("pretrained", (True, False))
    def test_pretrained_kw(self, pretrained):
371
        @handle_legacy_interface(weights=("pretrained", self.ModelWeights.Sentinel))
Philip Meier's avatar
Philip Meier committed
372
373
374
375
376
377
378
379
380
        def builder(*, weights=None):
            pass

        with pytest.warns(UserWarning, match="deprecated"):
            builder(pretrained)

    @pytest.mark.parametrize("pretrained", (True, False))
    @pytest.mark.parametrize("positional", (True, False))
    def test_equivalent_behavior_weights(self, pretrained, positional):
381
        @handle_legacy_interface(weights=("pretrained", self.ModelWeights.Sentinel))
Philip Meier's avatar
Philip Meier committed
382
383
384
385
        def builder(*, weights=None):
            pass

        args, kwargs = ((pretrained,), dict()) if positional else ((), dict(pretrained=pretrained))
386
        with pytest.warns(UserWarning, match=f"weights={self.ModelWeights.Sentinel if pretrained else None}"):
Philip Meier's avatar
Philip Meier committed
387
388
389
390
391
392
393
394
            builder(*args, **kwargs)

    def test_multi_params(self):
        weights_params = ("weights", "weights_other")
        pretrained_params = [param.replace("weights", "pretrained") for param in weights_params]

        @handle_legacy_interface(
            **{
395
                weights_param: (pretrained_param, self.ModelWeights.Sentinel)
Philip Meier's avatar
Philip Meier committed
396
397
398
399
400
401
402
403
404
405
406
407
408
409
                for weights_param, pretrained_param in zip(weights_params, pretrained_params)
            }
        )
        def builder(*, weights=None, weights_other=None):
            pass

        for pretrained_param in pretrained_params:
            with pytest.warns(UserWarning, match="deprecated"):
                builder(**{pretrained_param: True})

    def test_default_callable(self):
        @handle_legacy_interface(
            weights=(
                "pretrained",
410
                lambda kwargs: self.ModelWeights.Sentinel if kwargs["flag"] else None,
Philip Meier's avatar
Philip Meier committed
411
412
413
414
415
416
417
418
419
420
            )
        )
        def builder(*, weights=None, flag):
            pass

        with pytest.warns(UserWarning, match="deprecated"):
            builder(pretrained=True, flag=True)

        with pytest.raises(ValueError, match="weights"):
            builder(pretrained=True, flag=False)
421
422
423
424
425
426
427
428

    @pytest.mark.parametrize(
        "model_fn",
        [fn for fn in TM.list_model_fns(models) if fn.__name__ not in {"vit_h_14", "regnet_y_128gf"}]
        + TM.list_model_fns(models.detection)
        + TM.list_model_fns(models.quantization)
        + TM.list_model_fns(models.segmentation)
        + TM.list_model_fns(models.video)
429
430
431
432
433
        + TM.list_model_fns(models.optical_flow)
        + [
            lambda pretrained: resnet_fpn_backbone(backbone_name="resnet50", pretrained=pretrained),
            lambda pretrained: mobilenet_backbone(backbone_name="mobilenet_v2", fpn=False, pretrained=pretrained),
        ],
434
435
436
437
438
    )
    @run_if_test_with_extended
    def test_pretrained_deprecation(self, model_fn):
        with pytest.warns(UserWarning, match="deprecated"):
            model_fn(pretrained=True)