test_extended_models.py 15.9 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

    if model_fn.__name__ == "maskrcnn_resnet50_fpn_v2":
        pytest.skip(reason="FIXME https://github.com/pytorch/vision/issues/7349")

189
190
191
192
193
194
    # list of all possible supported high-level fields for weights meta-data
    permitted_fields = {
        "backend",
        "categories",
        "keypoint_names",
        "license",
195
        "_metrics",
196
        "min_size",
197
        "min_temporal_size",
198
199
200
        "num_params",
        "recipe",
        "unquantized",
201
        "_docs",
202
        "_ops",
Nicolas Hug's avatar
Nicolas Hug committed
203
        "_file_size",
204
205
    }
    # mandatory fields for each computer vision task
206
    classification_fields = {"categories", ("_metrics", "ImageNet-1K", "acc@1"), ("_metrics", "ImageNet-1K", "acc@5")}
207
    defaults = {
Nicolas Hug's avatar
Nicolas Hug committed
208
        "all": {"_metrics", "min_size", "num_params", "recipe", "_docs", "_file_size", "_ops"},
209
        "models": classification_fields,
210
        "detection": {"categories", ("_metrics", "COCO-val2017", "box_map")},
211
        "quantization": classification_fields | {"backend", "unquantized"},
212
213
214
215
216
217
        "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")},
218
        "optical_flow": set(),
219
    }
220
    model_name = model_fn.__name__
221
    module_name = model_fn.__module__.split(".")[-2]
222
    expected_fields = defaults["all"] | defaults[module_name]
223

224
    weights_enum = get_model_weights(model_fn)
225
226
    if len(weights_enum) == 0:
        pytest.skip(f"Model '{model_name}' doesn't have any pre-trained weights.")
227
228

    problematic_weights = {}
229
    incorrect_meta = []
230
    bad_names = []
231
    for w in weights_enum:
232
233
234
235
236
237
238
        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
239
240
241
        unsupported_fields = set(w.meta.keys()) - permitted_fields
        if missing_fields or unsupported_fields:
            problematic_weights[w] = {"missing": missing_fields, "unsupported": unsupported_fields}
242
243

        if w == weights_enum.DEFAULT or any(w.meta[k] != weights_enum.DEFAULT.meta[k] for k in ["num_params", "_ops"]):
244
            if module_name == "quantization":
245
                # parameters() count doesn't work well with quantization, so we check against the non-quantized
246
                unquantized_w = w.meta.get("unquantized")
247
248
249
250
251
252
253
254
255
                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"))

256
            else:
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
                # 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"))

273
274
        if not w.name.isupper():
            bad_names.append(w)
275

Nicolas Hug's avatar
Nicolas Hug committed
276
277
        if get_file_size_mb(w) != w.meta.get("_file_size"):
            incorrect_meta.append((w, "_file_size"))
278

279
    assert not problematic_weights
280
    assert not incorrect_meta
281
    assert not bad_names
282
283


284
@pytest.mark.parametrize(
285
    "model_fn",
286
287
288
289
290
291
    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),
292
)
293
294
295
@run_if_test_with_extended
def test_transforms_jit(model_fn):
    model_name = model_fn.__name__
296
    weights_enum = get_model_weights(model_fn)
297
298
299
    if len(weights_enum) == 0:
        pytest.skip(f"Model '{model_name}' doesn't have any pre-trained weights.")

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

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

332
333
334
335
336
337
338
    problematic_weights = []
    for w in weights_enum:
        transforms = w.transforms()
        try:
            TM._check_jit_scriptable(transforms, args)
        except Exception:
            problematic_weights.append(w)
339

340
    assert not problematic_weights
Philip Meier's avatar
Philip Meier committed
341
342
343
344
345


# With this filter, every unexpected warning will be turned into an error
@pytest.mark.filterwarnings("error")
class TestHandleLegacyInterface:
346
    class ModelWeights(WeightsEnum):
Philip Meier's avatar
Philip Meier committed
347
348
349
350
351
352
353
        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"),
354
            pytest.param(dict(weights=ModelWeights.Sentinel), id="Weights"),
Philip Meier's avatar
Philip Meier committed
355
356
357
        ],
    )
    def test_no_warn(self, kwargs):
358
        @handle_legacy_interface(weights=("pretrained", self.ModelWeights.Sentinel))
Philip Meier's avatar
Philip Meier committed
359
360
361
362
363
364
365
        def builder(*, weights=None):
            pass

        builder(**kwargs)

    @pytest.mark.parametrize("pretrained", (True, False))
    def test_pretrained_pos(self, pretrained):
366
        @handle_legacy_interface(weights=("pretrained", self.ModelWeights.Sentinel))
Philip Meier's avatar
Philip Meier committed
367
368
369
370
371
372
373
374
        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):
375
        @handle_legacy_interface(weights=("pretrained", self.ModelWeights.Sentinel))
Philip Meier's avatar
Philip Meier committed
376
377
378
379
380
381
382
383
384
        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):
385
        @handle_legacy_interface(weights=("pretrained", self.ModelWeights.Sentinel))
Philip Meier's avatar
Philip Meier committed
386
387
388
389
        def builder(*, weights=None):
            pass

        args, kwargs = ((pretrained,), dict()) if positional else ((), dict(pretrained=pretrained))
390
        with pytest.warns(UserWarning, match=f"weights={self.ModelWeights.Sentinel if pretrained else None}"):
Philip Meier's avatar
Philip Meier committed
391
392
393
394
395
396
397
398
            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(
            **{
399
                weights_param: (pretrained_param, self.ModelWeights.Sentinel)
Philip Meier's avatar
Philip Meier committed
400
401
402
403
404
405
406
407
408
409
410
411
412
413
                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",
414
                lambda kwargs: self.ModelWeights.Sentinel if kwargs["flag"] else None,
Philip Meier's avatar
Philip Meier committed
415
416
417
418
419
420
421
422
423
424
            )
        )
        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)
425
426
427
428
429
430
431
432

    @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)
433
434
435
436
437
        + 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),
        ],
438
439
440
441
442
    )
    @run_if_test_with_extended
    def test_pretrained_deprecation(self, model_fn):
        with pytest.warns(UserWarning, match="deprecated"):
            model_fn(pretrained=True)