test_models.py 25.8 KB
Newer Older
1
import functools
2
import io
3
4
import operator
import os
5
import sys
6
7
import traceback
import warnings
8
from collections import OrderedDict
9
10

import pytest
11
import torch
12
import torch.fx
13
import torch.nn as nn
14
import torchvision
15
16
from _utils_internal import get_relative_path
from common_utils import map_nested_tensor_object, freeze_rng_state, set_rng_seed, cpu_and_gpu, needs_cuda
17
from torchvision import models
18

eellison's avatar
eellison committed
19

20
ACCEPT = os.getenv("EXPECTTEST_ACCEPT", "0") == "1"
21
22


23
24
25
def get_available_classification_models():
    # TODO add a registration mechanism to torchvision.models
    return [k for k, v in models.__dict__.items() if callable(v) and k[0].lower() == k[0] and k[0] != "_"]
26
27
28
29
30


def get_available_segmentation_models():
    # TODO add a registration mechanism to torchvision.models
    return [k for k, v in models.segmentation.__dict__.items() if callable(v) and k[0].lower() == k[0] and k[0] != "_"]
31
32


33
34
35
36
37
def get_available_detection_models():
    # TODO add a registration mechanism to torchvision.models
    return [k for k, v in models.detection.__dict__.items() if callable(v) and k[0].lower() == k[0] and k[0] != "_"]


38
39
40
41
42
def get_available_video_models():
    # TODO add a registration mechanism to torchvision.models
    return [k for k, v in models.video.__dict__.items() if callable(v) and k[0].lower() == k[0] and k[0] != "_"]


43
44
45
46
47
def get_available_quantizable_models():
    # TODO add a registration mechanism to torchvision.models
    return [k for k, v in models.quantization.__dict__.items() if callable(v) and k[0].lower() == k[0] and k[0] != "_"]


48
49
50
51
52
53
def _get_expected_file(name=None):
    # Determine expected file based on environment
    expected_file_base = get_relative_path(os.path.realpath(__file__), "expect")

    # Note: for legacy reasons, the reference file names all had "ModelTest.test_" in their names
    # We hardcode it here to avoid having to re-generate the reference files
54
    expected_file = expected_file = os.path.join(expected_file_base, "ModelTester.test_" + name)
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
    expected_file += "_expect.pkl"

    if not ACCEPT and not os.path.exists(expected_file):
        raise RuntimeError(
            f"No expect file exists for {os.path.basename(expected_file)} in {expected_file}; "
            "to accept the current output, re-run the failing test after setting the EXPECTTEST_ACCEPT "
            "env variable. For example: EXPECTTEST_ACCEPT=1 pytest test/test_models.py -k alexnet"
        )

    return expected_file


def _assert_expected(output, name, prec):
    """Test that a python value matches the recorded contents of a file
    based on a "check" name. The value must be
    pickable with `torch.save`. This file
    is placed in the 'expect' directory in the same directory
    as the test script. You can automatically update the recorded test
    output using an EXPECTTEST_ACCEPT=1 env variable.
    """
    expected_file = _get_expected_file(name)

    if ACCEPT:
        filename = {os.path.basename(expected_file)}
        print("Accepting updated output for {}:\n\n{}".format(filename, output))
        torch.save(output, expected_file)
        MAX_PICKLE_SIZE = 50 * 1000  # 50 KB
        binary_size = os.path.getsize(expected_file)
        if binary_size > MAX_PICKLE_SIZE:
            raise RuntimeError("The output for {}, is larger than 50kb".format(filename))
    else:
        expected = torch.load(expected_file)
        rtol = atol = prec
        torch.testing.assert_close(output, expected, rtol=rtol, atol=atol, check_dtype=False)


def _check_jit_scriptable(nn_module, args, unwrapper=None, skip=False):
    """Check that a nn.Module's results in TorchScript match eager and that it can be exported"""

    def assert_export_import_module(m, args):
        """Check that the results of a model are the same after saving and loading"""
96

97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
        def get_export_import_copy(m):
            """Save and load a TorchScript model"""
            buffer = io.BytesIO()
            torch.jit.save(m, buffer)
            buffer.seek(0)
            imported = torch.jit.load(buffer)
            return imported

        m_import = get_export_import_copy(m)
        with freeze_rng_state():
            results = m(*args)
        with freeze_rng_state():
            results_from_imported = m_import(*args)
        tol = 3e-4
        try:
            torch.testing.assert_close(results, results_from_imported, atol=tol, rtol=tol)
113
        except ValueError:
114
115
116
117
118
119
            # custom check for the models that return named tuples:
            # we compare field by field while ignoring None as assert_close can't handle None
            for a, b in zip(results, results_from_imported):
                if a is not None:
                    torch.testing.assert_close(a, b, atol=tol, rtol=tol)

120
    TEST_WITH_SLOW = os.getenv("PYTORCH_TEST_WITH_SLOW", "0") == "1"
121
122
    if not TEST_WITH_SLOW or skip:
        # TorchScript is not enabled, skip these tests
123
124
125
126
127
128
129
130
        msg = (
            "The check_jit_scriptable test for {} was skipped. "
            "This test checks if the module's results in TorchScript "
            "match eager and that it can be exported. To run these "
            "tests make sure you set the environment variable "
            "PYTORCH_TEST_WITH_SLOW=1 and that the test is not "
            "manually skipped.".format(nn_module.__class__.__name__)
        )
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
        warnings.warn(msg, RuntimeWarning)
        return None

    sm = torch.jit.script(nn_module)

    with freeze_rng_state():
        eager_out = nn_module(*args)

    with freeze_rng_state():
        script_out = sm(*args)
        if unwrapper:
            script_out = unwrapper(script_out)

    torch.testing.assert_close(eager_out, script_out, atol=1e-4, rtol=1e-4)
    assert_export_import_module(sm, args)


148
149
150
151
152
153
154
def _check_fx_compatible(model, inputs):
    model_fx = torch.fx.symbolic_trace(model)
    out = model(inputs)
    out_fx = model_fx(inputs)
    torch.testing.assert_close(out, out_fx)


155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
def _check_input_backprop(model, inputs):
    if isinstance(inputs, list):
        requires_grad = list()
        for inp in inputs:
            requires_grad.append(inp.requires_grad)
            inp.requires_grad_(True)
    else:
        requires_grad = inputs.requires_grad
        inputs.requires_grad_(True)

    out = model(inputs)

    if isinstance(out, dict):
        out["out"].sum().backward()
    else:
        if isinstance(out[0], dict):
            out[0]["scores"].sum().backward()
        else:
            out[0].sum().backward()

    if isinstance(inputs, list):
        for i, inp in enumerate(inputs):
            assert inputs[i].grad is not None
            inp.requires_grad_(requires_grad[i])
    else:
        assert inputs.grad is not None
        inputs.requires_grad_(requires_grad)


184
185
186
# If 'unwrapper' is provided it will be called with the script model outputs
# before they are compared to the eager model outputs. This is useful if the
# model outputs are different between TorchScript / Eager mode
187
script_model_unwrapper = {
188
189
    "googlenet": lambda x: x.logits,
    "inception_v3": lambda x: x.logits,
190
    "fasterrcnn_resnet50_fpn": lambda x: x[1],
191
    "fasterrcnn_mobilenet_v3_large_fpn": lambda x: x[1],
192
    "fasterrcnn_mobilenet_v3_large_320_fpn": lambda x: x[1],
193
194
195
    "maskrcnn_resnet50_fpn": lambda x: x[1],
    "keypointrcnn_resnet50_fpn": lambda x: x[1],
    "retinanet_resnet50_fpn": lambda x: x[1],
196
    "ssd300_vgg16": lambda x: x[1],
197
    "ssdlite320_mobilenet_v3_large": lambda x: x[1],
198
}
199
200


201
202
203
204
205
206
207
208
209
210
211
212
213
214
# The following models exhibit flaky numerics under autocast in _test_*_model harnesses.
# This may be caused by the harness environment (e.g. num classes, input initialization
# via torch.rand), and does not prove autocast is unsuitable when training with real data
# (autocast has been used successfully with real data for some of these models).
# TODO:  investigate why autocast numerics are flaky in the harnesses.
#
# For the following models, _test_*_model harnesses skip numerical checks on outputs when
# trying autocast. However, they still try an autocasted forward pass, so they still ensure
# autocast coverage suffices to prevent dtype errors in each model.
autocast_flaky_numerics = (
    "inception_v3",
    "resnet101",
    "resnet152",
    "wide_resnet101_2",
215
216
    "deeplabv3_resnet50",
    "deeplabv3_resnet101",
217
    "deeplabv3_mobilenet_v3_large",
218
219
    "fcn_resnet50",
    "fcn_resnet101",
220
    "lraspp_mobilenet_v3_large",
221
    "maskrcnn_resnet50_fpn",
222
223
224
)


225
226
227
# The following contains configuration parameters for all models which are used by
# the _test_*_model methods.
_model_params = {
228
229
230
231
232
233
234
    "inception_v3": {"input_shape": (1, 3, 299, 299)},
    "retinanet_resnet50_fpn": {
        "num_classes": 20,
        "score_thresh": 0.01,
        "min_size": 224,
        "max_size": 224,
        "input_shape": (3, 224, 224),
235
    },
236
237
238
239
240
241
    "keypointrcnn_resnet50_fpn": {
        "num_classes": 2,
        "min_size": 224,
        "max_size": 224,
        "box_score_thresh": 0.15,
        "input_shape": (3, 224, 224),
242
    },
243
244
245
246
247
    "fasterrcnn_resnet50_fpn": {
        "num_classes": 20,
        "min_size": 224,
        "max_size": 224,
        "input_shape": (3, 224, 224),
248
    },
249
250
251
252
253
    "maskrcnn_resnet50_fpn": {
        "num_classes": 10,
        "min_size": 224,
        "max_size": 224,
        "input_shape": (3, 224, 224),
254
    },
255
256
    "fasterrcnn_mobilenet_v3_large_fpn": {
        "box_score_thresh": 0.02076,
257
    },
258
259
260
261
    "fasterrcnn_mobilenet_v3_large_320_fpn": {
        "box_score_thresh": 0.02076,
        "rpn_pre_nms_top_n_test": 1000,
        "rpn_post_nms_top_n_test": 1000,
262
263
264
265
    },
}


Anirudh's avatar
Anirudh committed
266
267
268
269
270
271
272
273
274
275
def _make_sliced_model(model, stop_layer):
    layers = OrderedDict()
    for name, layer in model.named_children():
        layers[name] = layer
        if name == stop_layer:
            break
    new_model = torch.nn.Sequential(layers)
    return new_model


276
@pytest.mark.parametrize("model_name", ["densenet121", "densenet169", "densenet201", "densenet161"])
Anirudh's avatar
Anirudh committed
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
def test_memory_efficient_densenet(model_name):
    input_shape = (1, 3, 300, 300)
    x = torch.rand(input_shape)

    model1 = models.__dict__[model_name](num_classes=50, memory_efficient=True)
    params = model1.state_dict()
    num_params = sum([x.numel() for x in model1.parameters()])
    model1.eval()
    out1 = model1(x)
    out1.sum().backward()
    num_grad = sum([x.grad.numel() for x in model1.parameters() if x.grad is not None])

    model2 = models.__dict__[model_name](num_classes=50, memory_efficient=False)
    model2.load_state_dict(params)
    model2.eval()
    out2 = model2(x)

    assert num_params == num_grad
    torch.testing.assert_close(out1, out2, rtol=0.0, atol=1e-5)

297
298
299
    _check_input_backprop(model1, x)
    _check_input_backprop(model2, x)

Anirudh's avatar
Anirudh committed
300

301
302
303
@pytest.mark.parametrize("dilate_layer_2", (True, False))
@pytest.mark.parametrize("dilate_layer_3", (True, False))
@pytest.mark.parametrize("dilate_layer_4", (True, False))
Anirudh's avatar
Anirudh committed
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
def test_resnet_dilation(dilate_layer_2, dilate_layer_3, dilate_layer_4):
    # TODO improve tests to also check that each layer has the right dimensionality
    model = models.__dict__["resnet50"](replace_stride_with_dilation=(dilate_layer_2, dilate_layer_3, dilate_layer_4))
    model = _make_sliced_model(model, stop_layer="layer4")
    model.eval()
    x = torch.rand(1, 3, 224, 224)
    out = model(x)
    f = 2 ** sum((dilate_layer_2, dilate_layer_3, dilate_layer_4))
    assert out.shape == (1, 2048, 7 * f, 7 * f)


def test_mobilenet_v2_residual_setting():
    model = models.__dict__["mobilenet_v2"](inverted_residual_setting=[[1, 16, 1, 1], [6, 24, 2, 2]])
    model.eval()
    x = torch.rand(1, 3, 224, 224)
    out = model(x)
    assert out.shape[-1] == 1000


323
@pytest.mark.parametrize("model_name", ["mobilenet_v2", "mobilenet_v3_large", "mobilenet_v3_small"])
Anirudh's avatar
Anirudh committed
324
325
326
327
328
329
330
331
def test_mobilenet_norm_layer(model_name):
    model = models.__dict__[model_name]()
    assert any(isinstance(x, nn.BatchNorm2d) for x in model.modules())

    def get_gn(num_channels):
        return nn.GroupNorm(32, num_channels)

    model = models.__dict__[model_name](norm_layer=get_gn)
332
    assert not (any(isinstance(x, nn.BatchNorm2d) for x in model.modules()))
Anirudh's avatar
Anirudh committed
333
334
335
336
337
338
    assert any(isinstance(x, nn.GroupNorm) for x in model.modules())


def test_inception_v3_eval():
    # replacement for models.inception_v3(pretrained=True) that does not download weights
    kwargs = {}
339
340
341
    kwargs["transform_input"] = True
    kwargs["aux_logits"] = True
    kwargs["init_weights"] = False
Anirudh's avatar
Anirudh committed
342
343
344
345
346
347
348
    name = "inception_v3"
    model = models.Inception3(**kwargs)
    model.aux_logits = False
    model.AuxLogits = None
    model = model.eval()
    x = torch.rand(1, 3, 299, 299)
    _check_jit_scriptable(model, (x,), unwrapper=script_model_unwrapper.get(name, None))
349
    _check_input_backprop(model, x)
Anirudh's avatar
Anirudh committed
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364


def test_fasterrcnn_double():
    model = models.detection.fasterrcnn_resnet50_fpn(num_classes=50, pretrained_backbone=False)
    model.double()
    model.eval()
    input_shape = (3, 300, 300)
    x = torch.rand(input_shape, dtype=torch.float64)
    model_input = [x]
    out = model(model_input)
    assert model_input[0] is x
    assert len(out) == 1
    assert "boxes" in out[0]
    assert "scores" in out[0]
    assert "labels" in out[0]
365
    _check_input_backprop(model, model_input)
Anirudh's avatar
Anirudh committed
366
367
368
369
370


def test_googlenet_eval():
    # replacement for models.googlenet(pretrained=True) that does not download weights
    kwargs = {}
371
372
373
    kwargs["transform_input"] = True
    kwargs["aux_logits"] = True
    kwargs["init_weights"] = False
Anirudh's avatar
Anirudh committed
374
375
376
377
378
379
380
381
    name = "googlenet"
    model = models.GoogLeNet(**kwargs)
    model.aux_logits = False
    model.aux1 = None
    model.aux2 = None
    model = model.eval()
    x = torch.rand(1, 3, 224, 224)
    _check_jit_scriptable(model, (x,), unwrapper=script_model_unwrapper.get(name, None))
382
    _check_input_backprop(model, x)
Anirudh's avatar
Anirudh committed
383
384
385
386
387
388
389
390
391
392
393
394
395
396


@needs_cuda
def test_fasterrcnn_switch_devices():
    def checkOut(out):
        assert len(out) == 1
        assert "boxes" in out[0]
        assert "scores" in out[0]
        assert "labels" in out[0]

    model = models.detection.fasterrcnn_resnet50_fpn(num_classes=50, pretrained_backbone=False)
    model.cuda()
    model.eval()
    input_shape = (3, 300, 300)
397
    x = torch.rand(input_shape, device="cuda")
Anirudh's avatar
Anirudh committed
398
399
400
401
402
403
404
    model_input = [x]
    out = model(model_input)
    assert model_input[0] is x

    checkOut(out)

    with torch.cuda.amp.autocast():
405
        out = model(model_input)
406

Anirudh's avatar
Anirudh committed
407
    checkOut(out)
408

409
410
    _check_input_backprop(model, model_input)

Anirudh's avatar
Anirudh committed
411
412
413
414
    # now switch to cpu and make sure it works
    model.cpu()
    x = x.cpu()
    out_cpu = model([x])
415

Anirudh's avatar
Anirudh committed
416
    checkOut(out_cpu)
417

418
419
    _check_input_backprop(model, [x])

420

Anirudh's avatar
Anirudh committed
421
def test_generalizedrcnn_transform_repr():
422

Anirudh's avatar
Anirudh committed
423
424
425
    min_size, max_size = 224, 299
    image_mean = [0.485, 0.456, 0.406]
    image_std = [0.229, 0.224, 0.225]
426

427
428
429
    t = models.detection.transform.GeneralizedRCNNTransform(
        min_size=min_size, max_size=max_size, image_mean=image_mean, image_std=image_std
    )
430

Anirudh's avatar
Anirudh committed
431
    # Check integrity of object __repr__ attribute
432
433
434
435
    expected_string = "GeneralizedRCNNTransform("
    _indent = "\n    "
    expected_string += "{0}Normalize(mean={1}, std={2})".format(_indent, image_mean, image_std)
    expected_string += "{0}Resize(min_size=({1},), max_size={2}, ".format(_indent, min_size, max_size)
Anirudh's avatar
Anirudh committed
436
437
    expected_string += "mode='bilinear')\n)"
    assert t.__repr__() == expected_string
438
439


440
441
@pytest.mark.parametrize("model_name", get_available_classification_models())
@pytest.mark.parametrize("dev", cpu_and_gpu())
442
def test_classification_model(model_name, dev):
Anirudh's avatar
Anirudh committed
443
444
    set_rng_seed(0)
    defaults = {
445
446
        "num_classes": 50,
        "input_shape": (1, 3, 224, 224),
Anirudh's avatar
Anirudh committed
447
448
    }
    kwargs = {**defaults, **_model_params.get(model_name, {})}
449
    input_shape = kwargs.pop("input_shape")
Anirudh's avatar
Anirudh committed
450
451
452
453
454
455
456
457
458

    model = models.__dict__[model_name](**kwargs)
    model.eval().to(device=dev)
    # RNG always on CPU, to ensure x in cuda tests is bitwise identical to x in cpu tests
    x = torch.rand(input_shape).to(device=dev)
    out = model(x)
    _assert_expected(out.cpu(), model_name, prec=0.1)
    assert out.shape[-1] == 50
    _check_jit_scriptable(model, (x,), unwrapper=script_model_unwrapper.get(model_name, None))
459
    _check_fx_compatible(model, x)
Anirudh's avatar
Anirudh committed
460
461
462
463
464
465
466
467

    if dev == torch.device("cuda"):
        with torch.cuda.amp.autocast():
            out = model(x)
            # See autocast_flaky_numerics comment at top of file.
            if model_name not in autocast_flaky_numerics:
                _assert_expected(out.cpu(), model_name, prec=0.1)
            assert out.shape[-1] == 50
468

469
470
    _check_input_backprop(model, x)

471

472
473
@pytest.mark.parametrize("model_name", get_available_segmentation_models())
@pytest.mark.parametrize("dev", cpu_and_gpu())
474
def test_segmentation_model(model_name, dev):
Anirudh's avatar
Anirudh committed
475
476
    set_rng_seed(0)
    defaults = {
477
478
479
        "num_classes": 10,
        "pretrained_backbone": False,
        "input_shape": (1, 3, 32, 32),
Anirudh's avatar
Anirudh committed
480
481
    }
    kwargs = {**defaults, **_model_params.get(model_name, {})}
482
    input_shape = kwargs.pop("input_shape")
Anirudh's avatar
Anirudh committed
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510

    model = models.segmentation.__dict__[model_name](**kwargs)
    model.eval().to(device=dev)
    # RNG always on CPU, to ensure x in cuda tests is bitwise identical to x in cpu tests
    x = torch.rand(input_shape).to(device=dev)
    out = model(x)["out"]

    def check_out(out):
        prec = 0.01
        try:
            # We first try to assert the entire output if possible. This is not
            # only the best way to assert results but also handles the cases
            # where we need to create a new expected result.
            _assert_expected(out.cpu(), model_name, prec=prec)
        except AssertionError:
            # Unfortunately some segmentation models are flaky with autocast
            # so instead of validating the probability scores, check that the class
            # predictions match.
            expected_file = _get_expected_file(model_name)
            expected = torch.load(expected_file)
            torch.testing.assert_close(out.argmax(dim=1), expected.argmax(dim=1), rtol=prec, atol=prec)
            return False  # Partial validation performed

        return True  # Full validation performed

    full_validation = check_out(out)

    _check_jit_scriptable(model, (x,), unwrapper=script_model_unwrapper.get(model_name, None))
511
    _check_fx_compatible(model, x)
Anirudh's avatar
Anirudh committed
512
513
514
515
516
517
518
519
520

    if dev == torch.device("cuda"):
        with torch.cuda.amp.autocast():
            out = model(x)["out"]
            # See autocast_flaky_numerics comment at top of file.
            if model_name not in autocast_flaky_numerics:
                full_validation &= check_out(out)

    if not full_validation:
521
522
523
524
525
526
        msg = (
            "The output of {} could only be partially validated. "
            "This is likely due to unit-test flakiness, but you may "
            "want to do additional manual checks if you made "
            "significant changes to the codebase.".format(test_segmentation_model.__name__)
        )
Anirudh's avatar
Anirudh committed
527
528
        warnings.warn(msg, RuntimeWarning)
        pytest.skip(msg)
529

530
531
    _check_input_backprop(model, x)

532

533
534
@pytest.mark.parametrize("model_name", get_available_detection_models())
@pytest.mark.parametrize("dev", cpu_and_gpu())
535
def test_detection_model(model_name, dev):
Anirudh's avatar
Anirudh committed
536
537
    set_rng_seed(0)
    defaults = {
538
539
540
        "num_classes": 50,
        "pretrained_backbone": False,
        "input_shape": (3, 300, 300),
Anirudh's avatar
Anirudh committed
541
542
    }
    kwargs = {**defaults, **_model_params.get(model_name, {})}
543
    input_shape = kwargs.pop("input_shape")
Anirudh's avatar
Anirudh committed
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570

    model = models.detection.__dict__[model_name](**kwargs)
    model.eval().to(device=dev)
    # RNG always on CPU, to ensure x in cuda tests is bitwise identical to x in cpu tests
    x = torch.rand(input_shape).to(device=dev)
    model_input = [x]
    out = model(model_input)
    assert model_input[0] is x

    def check_out(out):
        assert len(out) == 1

        def compact(tensor):
            size = tensor.size()
            elements_per_sample = functools.reduce(operator.mul, size[1:], 1)
            if elements_per_sample > 30:
                return compute_mean_std(tensor)
            else:
                return subsample_tensor(tensor)

        def subsample_tensor(tensor):
            num_elems = tensor.size(0)
            num_samples = 20
            if num_elems <= num_samples:
                return tensor

            ith_index = num_elems // num_samples
571
            return tensor[ith_index - 1 :: ith_index]
Anirudh's avatar
Anirudh committed
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593

        def compute_mean_std(tensor):
            # can't compute mean of integral tensor
            tensor = tensor.to(torch.double)
            mean = torch.mean(tensor)
            std = torch.std(tensor)
            return {"mean": mean, "std": std}

        output = map_nested_tensor_object(out, tensor_map_fn=compact)
        prec = 0.01
        try:
            # We first try to assert the entire output if possible. This is not
            # only the best way to assert results but also handles the cases
            # where we need to create a new expected result.
            _assert_expected(output, model_name, prec=prec)
        except AssertionError:
            # Unfortunately detection models are flaky due to the unstable sort
            # in NMS. If matching across all outputs fails, use the same approach
            # as in NMSTester.test_nms_cuda to see if this is caused by duplicate
            # scores.
            expected_file = _get_expected_file(model_name)
            expected = torch.load(expected_file)
594
595
596
            torch.testing.assert_close(
                output[0]["scores"], expected[0]["scores"], rtol=prec, atol=prec, check_device=False, check_dtype=False
            )
Anirudh's avatar
Anirudh committed
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616

            # Note: Fmassa proposed turning off NMS by adapting the threshold
            # and then using the Hungarian algorithm as in DETR to find the
            # best match between output and expected boxes and eliminate some
            # of the flakiness. Worth exploring.
            return False  # Partial validation performed

        return True  # Full validation performed

    full_validation = check_out(out)
    _check_jit_scriptable(model, ([x],), unwrapper=script_model_unwrapper.get(model_name, None))

    if dev == torch.device("cuda"):
        with torch.cuda.amp.autocast():
            out = model(model_input)
            # See autocast_flaky_numerics comment at top of file.
            if model_name not in autocast_flaky_numerics:
                full_validation &= check_out(out)

    if not full_validation:
617
618
619
620
621
622
        msg = (
            "The output of {} could only be partially validated. "
            "This is likely due to unit-test flakiness, but you may "
            "want to do additional manual checks if you made "
            "significant changes to the codebase.".format(test_detection_model.__name__)
        )
Anirudh's avatar
Anirudh committed
623
624
        warnings.warn(msg, RuntimeWarning)
        pytest.skip(msg)
625

626
627
    _check_input_backprop(model, model_input)

628

629
@pytest.mark.parametrize("model_name", get_available_detection_models())
630
def test_detection_model_validation(model_name):
Anirudh's avatar
Anirudh committed
631
632
633
634
635
636
637
638
639
640
    set_rng_seed(0)
    model = models.detection.__dict__[model_name](num_classes=50, pretrained_backbone=False)
    input_shape = (3, 300, 300)
    x = [torch.rand(input_shape)]

    # validate that targets are present in training
    with pytest.raises(ValueError):
        model(x)

    # validate type
641
    targets = [{"boxes": 0.0}]
Anirudh's avatar
Anirudh committed
642
643
644
645
646
    with pytest.raises(ValueError):
        model(x, targets=targets)

    # validate boxes shape
    for boxes in (torch.rand((4,)), torch.rand((1, 5))):
647
        targets = [{"boxes": boxes}]
Anirudh's avatar
Anirudh committed
648
649
650
651
652
        with pytest.raises(ValueError):
            model(x, targets=targets)

    # validate that no degenerate boxes are present
    boxes = torch.tensor([[1, 3, 1, 4], [2, 4, 3, 4]])
653
    targets = [{"boxes": boxes}]
Anirudh's avatar
Anirudh committed
654
655
    with pytest.raises(ValueError):
        model(x, targets=targets)
656

657

658
659
@pytest.mark.parametrize("model_name", get_available_video_models())
@pytest.mark.parametrize("dev", cpu_and_gpu())
660
def test_video_model(model_name, dev):
Anirudh's avatar
Anirudh committed
661
662
663
664
665
666
667
668
669
670
    # the default input shape is
    # bs * num_channels * clip_len * h *w
    input_shape = (1, 3, 4, 112, 112)
    # test both basicblock and Bottleneck
    model = models.video.__dict__[model_name](num_classes=50)
    model.eval().to(device=dev)
    # RNG always on CPU, to ensure x in cuda tests is bitwise identical to x in cpu tests
    x = torch.rand(input_shape).to(device=dev)
    out = model(x)
    _check_jit_scriptable(model, (x,), unwrapper=script_model_unwrapper.get(model_name, None))
671
    _check_fx_compatible(model, x)
Anirudh's avatar
Anirudh committed
672
673
674
675
676
677
    assert out.shape[-1] == 50

    if dev == torch.device("cuda"):
        with torch.cuda.amp.autocast():
            out = model(x)
            assert out.shape[-1] == 50
678

679
680
    _check_input_backprop(model, x)

681

682
683
684
685
686
687
688
689
@pytest.mark.skipif(
    not (
        "fbgemm" in torch.backends.quantized.supported_engines
        and "qnnpack" in torch.backends.quantized.supported_engines
    ),
    reason="This Pytorch Build has not been built with fbgemm and qnnpack",
)
@pytest.mark.parametrize("model_name", get_available_quantizable_models())
690
691
def test_quantized_classification_model(model_name):
    defaults = {
692
693
694
        "input_shape": (1, 3, 224, 224),
        "pretrained": False,
        "quantize": True,
695
696
    }
    kwargs = {**defaults, **_model_params.get(model_name, {})}
697
    input_shape = kwargs.pop("input_shape")
698
699
700
701
702
703

    # First check if quantize=True provides models that can run with input data
    model = torchvision.models.quantization.__dict__[model_name](**kwargs)
    x = torch.rand(input_shape)
    model(x)

704
    kwargs["quantize"] = False
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
    for eval_mode in [True, False]:
        model = torchvision.models.quantization.__dict__[model_name](**kwargs)
        if eval_mode:
            model.eval()
            model.qconfig = torch.quantization.default_qconfig
        else:
            model.train()
            model.qconfig = torch.quantization.default_qat_qconfig

        model.fuse_model()
        if eval_mode:
            torch.quantization.prepare(model, inplace=True)
        else:
            torch.quantization.prepare_qat(model, inplace=True)
            model.eval()

        torch.quantization.convert(model, inplace=True)

    try:
        torch.jit.script(model)
    except Exception as e:
        tb = traceback.format_exc()
        raise AssertionError(f"model cannot be scripted. Traceback = {str(tb)}") from e


730
if __name__ == "__main__":
731
    pytest.main([__file__])