test_extended_models.py 9.23 KB
Newer Older
1
import importlib
Philip Meier's avatar
Philip Meier committed
2
import os
3
4

import pytest
5
import test_models as TM
6
import torch
7
8
9
from torchvision import models
from torchvision.models._api import WeightsEnum, Weights
from torchvision.models._utils import handle_legacy_interface
10

11

12
13
14
15
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.",
)
16
17


18
19
20
21
22
23
def _get_parent_module(model_fn):
    parent_module_name = ".".join(model_fn.__module__.split(".")[:-1])
    module = importlib.import_module(parent_module_name)
    return module


24
25
26
27
28
29
30
31
32
33
34
35
36
def _get_model_weights(model_fn):
    module = _get_parent_module(model_fn)
    weights_name = "_QuantizedWeights" if module.__name__.split(".")[-1] == "quantization" else "_Weights"
    try:
        return next(
            v
            for k, v in module.__dict__.items()
            if k.endswith(weights_name) and k.replace(weights_name, "").lower() == model_fn.__name__
        )
    except StopIteration:
        return None


37
@pytest.mark.parametrize(
38
    "name, weight",
39
    [
40
41
        ("ResNet50_Weights.IMAGENET1K_V1", models.ResNet50_Weights.IMAGENET1K_V1),
        ("ResNet50_Weights.DEFAULT", models.ResNet50_Weights.IMAGENET1K_V2),
42
        (
43
44
            "ResNet50_QuantizedWeights.DEFAULT",
            models.quantization.ResNet50_QuantizedWeights.IMAGENET1K_FBGEMM_V2,
45
        ),
46
        (
47
48
            "ResNet50_QuantizedWeights.IMAGENET1K_FBGEMM_V1",
            models.quantization.ResNet50_QuantizedWeights.IMAGENET1K_FBGEMM_V1,
49
        ),
50
51
    ],
)
52
53
def test_get_weight(name, weight):
    assert models.get_weight(name) == weight
54
55


56
57
58
59
60
61
@pytest.mark.parametrize(
    "model_fn",
    TM.get_models_from_module(models)
    + TM.get_models_from_module(models.detection)
    + TM.get_models_from_module(models.quantization)
    + TM.get_models_from_module(models.segmentation)
62
63
    + TM.get_models_from_module(models.video)
    + TM.get_models_from_module(models.optical_flow),
64
65
)
def test_naming_conventions(model_fn):
66
67
    weights_enum = _get_model_weights(model_fn)
    assert weights_enum is not None
68
    assert len(weights_enum) == 0 or hasattr(weights_enum, "DEFAULT")
69
70


71
72
73
74
75
76
@pytest.mark.parametrize(
    "model_fn",
    TM.get_models_from_module(models)
    + TM.get_models_from_module(models.detection)
    + TM.get_models_from_module(models.quantization)
    + TM.get_models_from_module(models.segmentation)
77
78
    + TM.get_models_from_module(models.video)
    + TM.get_models_from_module(models.optical_flow),
79
)
80
@run_if_test_with_extended
81
def test_schema_meta_validation(model_fn):
82
83
    # TODO: add list of permitted fields
    classification_fields = ["categories", "acc@1", "acc@5"]
84
    defaults = {
85
        "all": ["recipe", "num_params", "min_size"],
86
87
        "models": classification_fields,
        "detection": ["categories", "map"],
88
        "quantization": classification_fields + ["backend", "unquantized"],
89
90
        "segmentation": ["categories", "mIoU", "acc"],
        "video": classification_fields,
91
        "optical_flow": [],
92
    }
93
    model_name = model_fn.__name__
94
95
96
97
    module_name = model_fn.__module__.split(".")[-2]
    fields = set(defaults["all"] + defaults[module_name])

    weights_enum = _get_model_weights(model_fn)
98
99
    if len(weights_enum) == 0:
        pytest.skip(f"Model '{model_name}' doesn't have any pre-trained weights.")
100
101

    problematic_weights = {}
102
    incorrect_params = []
103
    bad_names = []
104
105
106
107
    for w in weights_enum:
        missing_fields = fields - set(w.meta.keys())
        if missing_fields:
            problematic_weights[w] = missing_fields
108
        if w == weights_enum.DEFAULT:
109
            if module_name == "quantization":
110
                # parameters() count doesn't work well with quantization, so we check against the non-quantized
111
112
113
114
115
116
117
                unquantized_w = w.meta.get("unquantized")
                if unquantized_w is not None and w.meta.get("num_params") != unquantized_w.meta.get("num_params"):
                    incorrect_params.append(w)
            else:
                if w.meta.get("num_params") != sum(p.numel() for p in model_fn(weights=w).parameters()):
                    incorrect_params.append(w)
        else:
118
            if w.meta.get("num_params") != weights_enum.DEFAULT.meta.get("num_params"):
119
120
                if w.meta.get("num_params") != sum(p.numel() for p in model_fn(weights=w).parameters()):
                    incorrect_params.append(w)
121
122
        if not w.name.isupper():
            bad_names.append(w)
123
124

    assert not problematic_weights
125
    assert not incorrect_params
126
    assert not bad_names
127
128


129
@pytest.mark.parametrize(
130
131
132
133
134
    "model_fn",
    TM.get_models_from_module(models)
    + TM.get_models_from_module(models.detection)
    + TM.get_models_from_module(models.quantization)
    + TM.get_models_from_module(models.segmentation)
135
136
    + TM.get_models_from_module(models.video)
    + TM.get_models_from_module(models.optical_flow),
137
)
138
139
140
141
142
143
144
@run_if_test_with_extended
def test_transforms_jit(model_fn):
    model_name = model_fn.__name__
    weights_enum = _get_model_weights(model_fn)
    if len(weights_enum) == 0:
        pytest.skip(f"Model '{model_name}' doesn't have any pre-trained weights.")

145
    defaults = {
146
147
148
        "models": {
            "input_shape": (1, 3, 224, 224),
        },
149
150
151
        "detection": {
            "input_shape": (3, 300, 300),
        },
152
153
154
        "quantization": {
            "input_shape": (1, 3, 224, 224),
        },
155
156
157
158
        "segmentation": {
            "input_shape": (1, 3, 520, 520),
        },
        "video": {
159
            "input_shape": (1, 4, 112, 112, 3),
160
        },
161
162
163
        "optical_flow": {
            "input_shape": (1, 3, 128, 128),
        },
164
    }
165
    module_name = model_fn.__module__.split(".")[-2]
166

167
168
169
    kwargs = {**defaults[module_name], **TM._model_params.get(model_name, {})}
    input_shape = kwargs.pop("input_shape")
    x = torch.rand(input_shape)
170
    if module_name == "optical_flow":
171
        args = (x, x)
172
    else:
173
        args = (x,)
174

175
176
177
178
179
180
181
    problematic_weights = []
    for w in weights_enum:
        transforms = w.transforms()
        try:
            TM._check_jit_scriptable(transforms, args)
        except Exception:
            problematic_weights.append(w)
182

183
    assert not problematic_weights
Philip Meier's avatar
Philip Meier committed
184
185
186
187
188


# With this filter, every unexpected warning will be turned into an error
@pytest.mark.filterwarnings("error")
class TestHandleLegacyInterface:
189
    class ModelWeights(WeightsEnum):
Philip Meier's avatar
Philip Meier committed
190
191
192
193
194
195
196
        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"),
197
            pytest.param(dict(weights=ModelWeights.Sentinel), id="Weights"),
Philip Meier's avatar
Philip Meier committed
198
199
200
        ],
    )
    def test_no_warn(self, kwargs):
201
        @handle_legacy_interface(weights=("pretrained", self.ModelWeights.Sentinel))
Philip Meier's avatar
Philip Meier committed
202
203
204
205
206
207
208
        def builder(*, weights=None):
            pass

        builder(**kwargs)

    @pytest.mark.parametrize("pretrained", (True, False))
    def test_pretrained_pos(self, pretrained):
209
        @handle_legacy_interface(weights=("pretrained", self.ModelWeights.Sentinel))
Philip Meier's avatar
Philip Meier committed
210
211
212
213
214
215
216
217
        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):
218
        @handle_legacy_interface(weights=("pretrained", self.ModelWeights.Sentinel))
Philip Meier's avatar
Philip Meier committed
219
220
221
222
223
224
225
226
227
        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):
228
        @handle_legacy_interface(weights=("pretrained", self.ModelWeights.Sentinel))
Philip Meier's avatar
Philip Meier committed
229
230
231
232
        def builder(*, weights=None):
            pass

        args, kwargs = ((pretrained,), dict()) if positional else ((), dict(pretrained=pretrained))
233
        with pytest.warns(UserWarning, match=f"weights={self.ModelWeights.Sentinel if pretrained else None}"):
Philip Meier's avatar
Philip Meier committed
234
235
236
237
238
239
240
241
            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(
            **{
242
                weights_param: (pretrained_param, self.ModelWeights.Sentinel)
Philip Meier's avatar
Philip Meier committed
243
244
245
246
247
248
249
250
251
252
253
254
255
256
                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",
257
                lambda kwargs: self.ModelWeights.Sentinel if kwargs["flag"] else None,
Philip Meier's avatar
Philip Meier committed
258
259
260
261
262
263
264
265
266
267
            )
        )
        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)