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

import pytest
13
import torch
14
import torch.fx
15
import torch.nn as nn
16
17
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
18
from torchvision import models
19

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


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


32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
@pytest.fixture
def disable_weight_loading(mocker):
    """When testing models, the two slowest operations are the downloading of the weights to a file and loading them
    into the model. Unless, you want to test against specific weights, these steps can be disabled without any
    drawbacks.

    Including this fixture into the signature of your test, i.e. `test_foo(disable_weight_loading)`, will recurse
    through all models in `torchvision.models` and will patch all occurrences of the function
    `download_state_dict_from_url` as well as the method `load_state_dict` on all subclasses of `nn.Module` to be
    no-ops.

    .. warning:

        Loaded models are still executable as normal, but will always have random weights. Make sure to not use this
        fixture if you want to compare the model output against reference values.

    """
    starting_point = models
    function_name = "load_state_dict_from_url"
    method_name = "load_state_dict"

    module_names = {info.name for info in pkgutil.walk_packages(starting_point.__path__, f"{starting_point.__name__}.")}
    targets = {f"torchvision._internally_replaced_utils.{function_name}", f"torch.nn.Module.{method_name}"}
    for name in module_names:
        module = sys.modules.get(name)
        if not module:
            continue

        if function_name in module.__dict__:
            targets.add(f"{module.__name__}.{function_name}")

        targets.update(
            {
                f"{module.__name__}.{obj.__name__}.{method_name}"
                for obj in module.__dict__.values()
                if isinstance(obj, type) and issubclass(obj, nn.Module) and method_name in obj.__dict__
            }
        )

    for target in targets:
        # See https://github.com/pytorch/vision/pull/4867#discussion_r743677802 for details
        with contextlib.suppress(AttributeError):
            mocker.patch(target)


77
78
79
80
81
82
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
83
    expected_file = expected_file = os.path.join(expected_file_base, "ModelTester.test_" + name)
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


96
def _assert_expected(output, name, prec=None, atol=None, rtol=None):
97
98
99
100
101
102
103
104
105
106
107
    """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)}
108
        print(f"Accepting updated output for {filename}:\n\n{output}")
109
110
111
112
        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:
113
            raise RuntimeError(f"The output for {filename}, is larger than 50kb - got {binary_size}kb")
114
115
    else:
        expected = torch.load(expected_file)
116
117
        rtol = rtol or prec  # keeping prec param for legacy reason, but could be removed ideally
        atol = atol or prec
118
119
120
121
122
123
124
125
        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"""
126

127
128
129
130
131
132
133
134
135
136
137
138
139
140
        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
141
        torch.testing.assert_close(results, results_from_imported, atol=tol, rtol=tol)
142

143
    TEST_WITH_SLOW = os.getenv("PYTORCH_TEST_WITH_SLOW", "0") == "1"
144
145
    if not TEST_WITH_SLOW or skip:
        # TorchScript is not enabled, skip these tests
146
        msg = (
147
            f"The check_jit_scriptable test for {nn_module.__class__.__name__} was skipped. "
148
149
150
151
            "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 "
152
            "manually skipped."
153
        )
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
        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)


171
172
173
174
175
176
177
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)


178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
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)


207
208
209
# 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
210
script_model_unwrapper = {
211
212
    "googlenet": lambda x: x.logits,
    "inception_v3": lambda x: x.logits,
213
    "fasterrcnn_resnet50_fpn": lambda x: x[1],
214
    "fasterrcnn_mobilenet_v3_large_fpn": lambda x: x[1],
215
    "fasterrcnn_mobilenet_v3_large_320_fpn": lambda x: x[1],
216
217
218
    "maskrcnn_resnet50_fpn": lambda x: x[1],
    "keypointrcnn_resnet50_fpn": lambda x: x[1],
    "retinanet_resnet50_fpn": lambda x: x[1],
219
    "ssd300_vgg16": lambda x: x[1],
220
    "ssdlite320_mobilenet_v3_large": lambda x: x[1],
221
}
222
223


224
225
226
227
228
229
230
231
232
233
234
235
236
237
# 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",
238
239
    "deeplabv3_resnet50",
    "deeplabv3_resnet101",
240
    "deeplabv3_mobilenet_v3_large",
241
242
    "fcn_resnet50",
    "fcn_resnet101",
243
    "lraspp_mobilenet_v3_large",
244
    "maskrcnn_resnet50_fpn",
245
246
)

247
248
249
# The tests for the following quantized models are flaky possibly due to inconsistent
# rounding errors in different platforms. For this reason the input/output consistency
# tests under test_quantized_classification_model will be skipped for the following models.
250
quantized_flaky_models = ("inception_v3", "resnet50")
251

252

253
254
255
# The following contains configuration parameters for all models which are used by
# the _test_*_model methods.
_model_params = {
256
257
258
259
260
261
262
    "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),
263
    },
264
265
266
267
268
269
    "keypointrcnn_resnet50_fpn": {
        "num_classes": 2,
        "min_size": 224,
        "max_size": 224,
        "box_score_thresh": 0.15,
        "input_shape": (3, 224, 224),
270
    },
271
272
273
274
275
    "fasterrcnn_resnet50_fpn": {
        "num_classes": 20,
        "min_size": 224,
        "max_size": 224,
        "input_shape": (3, 224, 224),
276
    },
277
278
279
280
281
    "maskrcnn_resnet50_fpn": {
        "num_classes": 10,
        "min_size": 224,
        "max_size": 224,
        "input_shape": (3, 224, 224),
282
    },
283
284
    "fasterrcnn_mobilenet_v3_large_fpn": {
        "box_score_thresh": 0.02076,
285
    },
286
287
288
289
    "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,
290
291
292
293
    },
}


294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
# The following contains configuration and expected values to be used tests that are model specific
_model_tests_values = {
    "retinanet_resnet50_fpn": {
        "max_trainable": 5,
        "n_trn_params_per_layer": [36, 46, 65, 78, 88, 89],
    },
    "keypointrcnn_resnet50_fpn": {
        "max_trainable": 5,
        "n_trn_params_per_layer": [48, 58, 77, 90, 100, 101],
    },
    "fasterrcnn_resnet50_fpn": {
        "max_trainable": 5,
        "n_trn_params_per_layer": [30, 40, 59, 72, 82, 83],
    },
    "maskrcnn_resnet50_fpn": {
        "max_trainable": 5,
        "n_trn_params_per_layer": [42, 52, 71, 84, 94, 95],
    },
    "fasterrcnn_mobilenet_v3_large_fpn": {
        "max_trainable": 6,
        "n_trn_params_per_layer": [22, 23, 44, 70, 91, 97, 100],
    },
    "fasterrcnn_mobilenet_v3_large_320_fpn": {
        "max_trainable": 6,
        "n_trn_params_per_layer": [22, 23, 44, 70, 91, 97, 100],
    },
    "ssd300_vgg16": {
        "max_trainable": 5,
        "n_trn_params_per_layer": [45, 51, 57, 63, 67, 71],
    },
    "ssdlite320_mobilenet_v3_large": {
        "max_trainable": 6,
        "n_trn_params_per_layer": [96, 99, 138, 200, 239, 257, 266],
    },
}


Anirudh's avatar
Anirudh committed
331
332
333
334
335
336
337
338
339
340
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


341
342
@pytest.mark.parametrize("model_fn", [models.densenet121, models.densenet169, models.densenet201, models.densenet161])
def test_memory_efficient_densenet(model_fn):
Anirudh's avatar
Anirudh committed
343
344
345
    input_shape = (1, 3, 300, 300)
    x = torch.rand(input_shape)

346
    model1 = model_fn(num_classes=50, memory_efficient=True)
Anirudh's avatar
Anirudh committed
347
    params = model1.state_dict()
348
    num_params = sum(x.numel() for x in model1.parameters())
Anirudh's avatar
Anirudh committed
349
350
351
    model1.eval()
    out1 = model1(x)
    out1.sum().backward()
352
    num_grad = sum(x.grad.numel() for x in model1.parameters() if x.grad is not None)
Anirudh's avatar
Anirudh committed
353

354
    model2 = model_fn(num_classes=50, memory_efficient=False)
Anirudh's avatar
Anirudh committed
355
356
357
358
359
360
361
    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)

362
363
364
    _check_input_backprop(model1, x)
    _check_input_backprop(model2, x)

Anirudh's avatar
Anirudh committed
365

366
367
368
@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
369
370
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
371
    model = models.resnet50(replace_stride_with_dilation=(dilate_layer_2, dilate_layer_3, dilate_layer_4))
Anirudh's avatar
Anirudh committed
372
373
374
375
376
377
378
379
380
    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():
381
    model = models.mobilenet_v2(inverted_residual_setting=[[1, 16, 1, 1], [6, 24, 2, 2]])
Anirudh's avatar
Anirudh committed
382
383
384
385
386
387
    model.eval()
    x = torch.rand(1, 3, 224, 224)
    out = model(x)
    assert out.shape[-1] == 1000


388
389
390
@pytest.mark.parametrize("model_fn", [models.mobilenet_v2, models.mobilenet_v3_large, models.mobilenet_v3_small])
def test_mobilenet_norm_layer(model_fn):
    model = model_fn()
Anirudh's avatar
Anirudh committed
391
392
393
394
395
    assert any(isinstance(x, nn.BatchNorm2d) for x in model.modules())

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

396
    model = model_fn(norm_layer=get_gn)
397
    assert not (any(isinstance(x, nn.BatchNorm2d) for x in model.modules()))
Anirudh's avatar
Anirudh committed
398
399
400
401
402
403
    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 = {}
404
405
406
    kwargs["transform_input"] = True
    kwargs["aux_logits"] = True
    kwargs["init_weights"] = False
Anirudh's avatar
Anirudh committed
407
408
409
410
411
412
413
    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))
414
    _check_input_backprop(model, x)
Anirudh's avatar
Anirudh committed
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429


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]
430
    _check_input_backprop(model, model_input)
Anirudh's avatar
Anirudh committed
431
432
433
434
435


def test_googlenet_eval():
    # replacement for models.googlenet(pretrained=True) that does not download weights
    kwargs = {}
436
437
438
    kwargs["transform_input"] = True
    kwargs["aux_logits"] = True
    kwargs["init_weights"] = False
Anirudh's avatar
Anirudh committed
439
440
441
442
443
444
445
446
    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))
447
    _check_input_backprop(model, x)
Anirudh's avatar
Anirudh committed
448
449
450
451
452
453
454
455
456
457
458
459
460
461


@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)
462
    x = torch.rand(input_shape, device="cuda")
Anirudh's avatar
Anirudh committed
463
464
465
466
467
468
469
    model_input = [x]
    out = model(model_input)
    assert model_input[0] is x

    checkOut(out)

    with torch.cuda.amp.autocast():
470
        out = model(model_input)
471

Anirudh's avatar
Anirudh committed
472
    checkOut(out)
473

474
475
    _check_input_backprop(model, model_input)

Anirudh's avatar
Anirudh committed
476
477
478
479
    # now switch to cpu and make sure it works
    model.cpu()
    x = x.cpu()
    out_cpu = model([x])
480

Anirudh's avatar
Anirudh committed
481
    checkOut(out_cpu)
482

483
484
    _check_input_backprop(model, [x])

485

Anirudh's avatar
Anirudh committed
486
def test_generalizedrcnn_transform_repr():
487

Anirudh's avatar
Anirudh committed
488
489
490
    min_size, max_size = 224, 299
    image_mean = [0.485, 0.456, 0.406]
    image_std = [0.229, 0.224, 0.225]
491

492
493
494
    t = models.detection.transform.GeneralizedRCNNTransform(
        min_size=min_size, max_size=max_size, image_mean=image_mean, image_std=image_std
    )
495

Anirudh's avatar
Anirudh committed
496
    # Check integrity of object __repr__ attribute
497
498
    expected_string = "GeneralizedRCNNTransform("
    _indent = "\n    "
499
500
    expected_string += f"{_indent}Normalize(mean={image_mean}, std={image_std})"
    expected_string += f"{_indent}Resize(min_size=({min_size},), max_size={max_size}, "
Anirudh's avatar
Anirudh committed
501
502
    expected_string += "mode='bilinear')\n)"
    assert t.__repr__() == expected_string
503
504


505
@pytest.mark.parametrize("model_fn", get_models_from_module(models))
506
@pytest.mark.parametrize("dev", cpu_and_gpu())
507
def test_classification_model(model_fn, dev):
Anirudh's avatar
Anirudh committed
508
509
    set_rng_seed(0)
    defaults = {
510
511
        "num_classes": 50,
        "input_shape": (1, 3, 224, 224),
Anirudh's avatar
Anirudh committed
512
    }
513
    model_name = model_fn.__name__
Anirudh's avatar
Anirudh committed
514
    kwargs = {**defaults, **_model_params.get(model_name, {})}
515
    num_classes = kwargs.get("num_classes")
516
    input_shape = kwargs.pop("input_shape")
Anirudh's avatar
Anirudh committed
517

518
    model = model_fn(**kwargs)
Anirudh's avatar
Anirudh committed
519
520
521
522
523
    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)
524
    assert out.shape[-1] == num_classes
Anirudh's avatar
Anirudh committed
525
    _check_jit_scriptable(model, (x,), unwrapper=script_model_unwrapper.get(model_name, None))
526
    _check_fx_compatible(model, x)
Anirudh's avatar
Anirudh committed
527
528
529
530
531
532
533
534

    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
535

536
537
    _check_input_backprop(model, x)

538

539
@pytest.mark.parametrize("model_fn", get_models_from_module(models.segmentation))
540
@pytest.mark.parametrize("dev", cpu_and_gpu())
541
def test_segmentation_model(model_fn, dev):
Anirudh's avatar
Anirudh committed
542
543
    set_rng_seed(0)
    defaults = {
544
545
546
        "num_classes": 10,
        "pretrained_backbone": False,
        "input_shape": (1, 3, 32, 32),
Anirudh's avatar
Anirudh committed
547
    }
548
    model_name = model_fn.__name__
Anirudh's avatar
Anirudh committed
549
    kwargs = {**defaults, **_model_params.get(model_name, {})}
550
    input_shape = kwargs.pop("input_shape")
Anirudh's avatar
Anirudh committed
551

552
    model = model_fn(**kwargs)
Anirudh's avatar
Anirudh committed
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
    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))
579
    _check_fx_compatible(model, x)
Anirudh's avatar
Anirudh committed
580
581
582
583
584
585
586
587
588

    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:
589
        msg = (
590
            f"The output of {test_segmentation_model.__name__} could only be partially validated. "
591
592
            "This is likely due to unit-test flakiness, but you may "
            "want to do additional manual checks if you made "
593
            "significant changes to the codebase."
594
        )
Anirudh's avatar
Anirudh committed
595
596
        warnings.warn(msg, RuntimeWarning)
        pytest.skip(msg)
597

598
599
    _check_input_backprop(model, x)

600

601
@pytest.mark.parametrize("model_fn", get_models_from_module(models.detection))
602
@pytest.mark.parametrize("dev", cpu_and_gpu())
603
def test_detection_model(model_fn, dev):
Anirudh's avatar
Anirudh committed
604
605
    set_rng_seed(0)
    defaults = {
606
607
608
        "num_classes": 50,
        "pretrained_backbone": False,
        "input_shape": (3, 300, 300),
Anirudh's avatar
Anirudh committed
609
    }
610
    model_name = model_fn.__name__
Anirudh's avatar
Anirudh committed
611
    kwargs = {**defaults, **_model_params.get(model_name, {})}
612
    input_shape = kwargs.pop("input_shape")
Anirudh's avatar
Anirudh committed
613

614
    model = model_fn(**kwargs)
Anirudh's avatar
Anirudh committed
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
    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
640
            return tensor[ith_index - 1 :: ith_index]
Anirudh's avatar
Anirudh committed
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662

        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)
663
664
665
            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
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685

            # 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:
686
        msg = (
687
            f"The output of {test_detection_model.__name__} could only be partially validated. "
688
689
            "This is likely due to unit-test flakiness, but you may "
            "want to do additional manual checks if you made "
690
            "significant changes to the codebase."
691
        )
Anirudh's avatar
Anirudh committed
692
693
        warnings.warn(msg, RuntimeWarning)
        pytest.skip(msg)
694

695
696
    _check_input_backprop(model, model_input)

697

698
699
@pytest.mark.parametrize("model_fn", get_models_from_module(models.detection))
def test_detection_model_validation(model_fn):
Anirudh's avatar
Anirudh committed
700
    set_rng_seed(0)
701
    model = model_fn(num_classes=50, pretrained_backbone=False)
Anirudh's avatar
Anirudh committed
702
703
704
705
706
707
708
709
    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
710
    targets = [{"boxes": 0.0}]
Anirudh's avatar
Anirudh committed
711
712
713
714
715
    with pytest.raises(ValueError):
        model(x, targets=targets)

    # validate boxes shape
    for boxes in (torch.rand((4,)), torch.rand((1, 5))):
716
        targets = [{"boxes": boxes}]
Anirudh's avatar
Anirudh committed
717
718
719
720
721
        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]])
722
    targets = [{"boxes": boxes}]
Anirudh's avatar
Anirudh committed
723
724
    with pytest.raises(ValueError):
        model(x, targets=targets)
725

726

727
@pytest.mark.parametrize("model_fn", get_models_from_module(models.video))
728
@pytest.mark.parametrize("dev", cpu_and_gpu())
729
def test_video_model(model_fn, dev):
Anirudh's avatar
Anirudh committed
730
731
732
    # the default input shape is
    # bs * num_channels * clip_len * h *w
    input_shape = (1, 3, 4, 112, 112)
733
    model_name = model_fn.__name__
Anirudh's avatar
Anirudh committed
734
    # test both basicblock and Bottleneck
735
    model = model_fn(num_classes=50)
Anirudh's avatar
Anirudh committed
736
737
738
739
740
    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))
741
    _check_fx_compatible(model, x)
Anirudh's avatar
Anirudh committed
742
743
744
745
746
747
    assert out.shape[-1] == 50

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

749
750
    _check_input_backprop(model, x)

751

752
753
754
755
756
757
758
@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",
)
759
760
@pytest.mark.parametrize("model_fn", get_models_from_module(models.quantization))
def test_quantized_classification_model(model_fn):
761
    set_rng_seed(0)
762
    defaults = {
763
        "num_classes": 5,
764
765
766
        "input_shape": (1, 3, 224, 224),
        "pretrained": False,
        "quantize": True,
767
    }
768
    model_name = model_fn.__name__
769
    kwargs = {**defaults, **_model_params.get(model_name, {})}
770
    input_shape = kwargs.pop("input_shape")
771
772

    # First check if quantize=True provides models that can run with input data
773
    model = model_fn(**kwargs)
774
    model.eval()
775
    x = torch.rand(input_shape)
776
777
778
779
780
781
782
    out = model(x)

    if model_name not in quantized_flaky_models:
        _assert_expected(out, model_name + "_quantized", prec=0.1)
        assert out.shape[-1] == 5
        _check_jit_scriptable(model, (x,), unwrapper=script_model_unwrapper.get(model_name, None))
        _check_fx_compatible(model, x)
783

784
    kwargs["quantize"] = False
785
    for eval_mode in [True, False]:
786
        model = model_fn(**kwargs)
787
788
        if eval_mode:
            model.eval()
789
            model.qconfig = torch.ao.quantization.default_qconfig
790
791
        else:
            model.train()
792
            model.qconfig = torch.ao.quantization.default_qat_qconfig
793
794
795

        model.fuse_model()
        if eval_mode:
796
            torch.ao.quantization.prepare(model, inplace=True)
797
        else:
798
            torch.ao.quantization.prepare_qat(model, inplace=True)
799
800
            model.eval()

801
        torch.ao.quantization.convert(model, inplace=True)
802
803
804
805
806
807
808
809

    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


810
@pytest.mark.parametrize("model_fn", get_models_from_module(models.detection))
811
def test_detection_model_trainable_backbone_layers(model_fn, disable_weight_loading):
812
    model_name = model_fn.__name__
813
814
815
    max_trainable = _model_tests_values[model_name]["max_trainable"]
    n_trainable_params = []
    for trainable_layers in range(0, max_trainable + 1):
816
        model = model_fn(pretrained=False, pretrained_backbone=True, trainable_backbone_layers=trainable_layers)
817
818
819
820
821

        n_trainable_params.append(len([p for p in model.parameters() if p.requires_grad]))
    assert n_trainable_params == _model_tests_values[model_name]["n_trn_params_per_layer"]


822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
@needs_cuda
@pytest.mark.parametrize("model_builder", (models.optical_flow.raft_large, models.optical_flow.raft_small))
@pytest.mark.parametrize("scripted", (False, True))
def test_raft(model_builder, scripted):

    torch.manual_seed(0)

    # We need very small images, otherwise the pickle size would exceed the 50KB
    # As a resut we need to override the correlation pyramid to not downsample
    # too much, otherwise we would get nan values (effective H and W would be
    # reduced to 1)
    corr_block = models.optical_flow.raft.CorrBlock(num_levels=2, radius=2)

    model = model_builder(corr_block=corr_block).eval().to("cuda")
    if scripted:
        model = torch.jit.script(model)

    bs = 1
    img1 = torch.rand(bs, 3, 80, 72).cuda()
    img2 = torch.rand(bs, 3, 80, 72).cuda()

    preds = model(img1, img2)
    flow_pred = preds[-1]
    # Tolerance is fairly high, but there are 2 * H * W outputs to check
    # The .pkl were generated on the AWS cluter, on the CI it looks like the resuts are slightly different
    _assert_expected(flow_pred, name=model_builder.__name__, atol=1e-2, rtol=1)


850
if __name__ == "__main__":
851
    pytest.main([__file__])