test_models.py 31.6 KB
Newer Older
1
import contextlib
2
3
4
import functools
import operator
import os
5
6
import pkgutil
import sys
7
8
import traceback
import warnings
9
from collections import OrderedDict
10
from tempfile import TemporaryDirectory
11
from typing import Any
12
13

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

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


24
def get_models_from_module(module):
25
    # TODO add a registration mechanism to torchvision.models
26
27
28
29
30
    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"
    ]
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
77
@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)


78
79
80
81
82
83
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
84
    expected_file = expected_file = os.path.join(expected_file_base, "ModelTester.test_" + name)
85
86
87
88
89
90
91
92
93
94
95
96
    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


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

128
129
        def get_export_import_copy(m):
            """Save and load a TorchScript model"""
130
131
132
133
            with TemporaryDirectory() as dir:
                path = os.path.join(dir, "script.pt")
                m.save(path)
                imported = torch.jit.load(path)
134
135
136
            return imported

        m_import = get_export_import_copy(m)
137
        with torch.no_grad(), freeze_rng_state():
138
            results = m(*args)
139
        with torch.no_grad(), freeze_rng_state():
140
141
            results_from_imported = m_import(*args)
        tol = 3e-4
142
        torch.testing.assert_close(results, results_from_imported, atol=tol, rtol=tol)
143

144
    TEST_WITH_SLOW = os.getenv("PYTORCH_TEST_WITH_SLOW", "0") == "1"
145
146
    if not TEST_WITH_SLOW or skip:
        # TorchScript is not enabled, skip these tests
147
        msg = (
148
            f"The check_jit_scriptable test for {nn_module.__class__.__name__} was skipped. "
149
150
151
152
            "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 "
153
            "manually skipped."
154
        )
155
156
157
158
159
        warnings.warn(msg, RuntimeWarning)
        return None

    sm = torch.jit.script(nn_module)

160
    with torch.no_grad(), freeze_rng_state():
161
162
        eager_out = nn_module(*args)

163
    with torch.no_grad(), freeze_rng_state():
164
165
166
167
168
169
170
171
        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)


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


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
207
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)


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


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

249
250
251
# 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.
252
quantized_flaky_models = ("inception_v3", "resnet50")
253

254

255
256
257
# The following contains configuration parameters for all models which are used by
# the _test_*_model methods.
_model_params = {
258
259
260
261
262
263
264
    "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),
265
    },
266
267
268
269
270
271
    "keypointrcnn_resnet50_fpn": {
        "num_classes": 2,
        "min_size": 224,
        "max_size": 224,
        "box_score_thresh": 0.15,
        "input_shape": (3, 224, 224),
272
    },
273
274
275
276
277
    "fasterrcnn_resnet50_fpn": {
        "num_classes": 20,
        "min_size": 224,
        "max_size": 224,
        "input_shape": (3, 224, 224),
278
    },
Hu Ye's avatar
Hu Ye committed
279
280
281
282
283
284
285
    "fcos_resnet50_fpn": {
        "num_classes": 2,
        "score_thresh": 0.05,
        "min_size": 224,
        "max_size": 224,
        "input_shape": (3, 224, 224),
    },
286
287
288
289
290
    "maskrcnn_resnet50_fpn": {
        "num_classes": 10,
        "min_size": 224,
        "max_size": 224,
        "input_shape": (3, 224, 224),
291
    },
292
293
    "fasterrcnn_mobilenet_v3_large_fpn": {
        "box_score_thresh": 0.02076,
294
    },
295
296
297
298
    "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,
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
331
332
333
334
335
336
# 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],
    },
Hu Ye's avatar
Hu Ye committed
337
338
339
340
    "fcos_resnet50_fpn": {
        "max_trainable": 5,
        "n_trn_params_per_layer": [54, 64, 83, 96, 106, 107],
    },
341
342
343
}


Anirudh's avatar
Anirudh committed
344
345
346
347
348
349
350
351
352
353
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


354
355
@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
356
357
358
    input_shape = (1, 3, 300, 300)
    x = torch.rand(input_shape)

359
    model1 = model_fn(num_classes=50, memory_efficient=True)
Anirudh's avatar
Anirudh committed
360
    params = model1.state_dict()
361
    num_params = sum(x.numel() for x in model1.parameters())
Anirudh's avatar
Anirudh committed
362
363
364
    model1.eval()
    out1 = model1(x)
    out1.sum().backward()
365
    num_grad = sum(x.grad.numel() for x in model1.parameters() if x.grad is not None)
Anirudh's avatar
Anirudh committed
366

367
    model2 = model_fn(num_classes=50, memory_efficient=False)
Anirudh's avatar
Anirudh committed
368
369
370
371
372
373
374
    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)

375
376
377
    _check_input_backprop(model1, x)
    _check_input_backprop(model2, x)

Anirudh's avatar
Anirudh committed
378

379
380
381
@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
382
383
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
384
    model = models.resnet50(replace_stride_with_dilation=(dilate_layer_2, dilate_layer_3, dilate_layer_4))
Anirudh's avatar
Anirudh committed
385
386
387
388
389
390
391
392
393
    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():
394
    model = models.mobilenet_v2(inverted_residual_setting=[[1, 16, 1, 1], [6, 24, 2, 2]])
Anirudh's avatar
Anirudh committed
395
396
397
398
399
400
    model.eval()
    x = torch.rand(1, 3, 224, 224)
    out = model(x)
    assert out.shape[-1] == 1000


401
402
403
@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
404
405
406
407
408
    assert any(isinstance(x, nn.BatchNorm2d) for x in model.modules())

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

409
    model = model_fn(norm_layer=get_gn)
410
    assert not (any(isinstance(x, nn.BatchNorm2d) for x in model.modules()))
Anirudh's avatar
Anirudh committed
411
412
413
414
415
416
    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 = {}
417
418
419
    kwargs["transform_input"] = True
    kwargs["aux_logits"] = True
    kwargs["init_weights"] = False
Anirudh's avatar
Anirudh committed
420
421
422
423
424
425
426
    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))
427
    _check_input_backprop(model, x)
Anirudh's avatar
Anirudh committed
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442


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]
443
    _check_input_backprop(model, model_input)
Anirudh's avatar
Anirudh committed
444
445
446
447
448


def test_googlenet_eval():
    # replacement for models.googlenet(pretrained=True) that does not download weights
    kwargs = {}
449
450
451
    kwargs["transform_input"] = True
    kwargs["aux_logits"] = True
    kwargs["init_weights"] = False
Anirudh's avatar
Anirudh committed
452
453
454
455
456
457
458
459
    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))
460
    _check_input_backprop(model, x)
Anirudh's avatar
Anirudh committed
461
462
463
464
465
466
467
468
469
470
471
472
473
474


@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)
475
    x = torch.rand(input_shape, device="cuda")
Anirudh's avatar
Anirudh committed
476
477
478
479
480
481
482
    model_input = [x]
    out = model(model_input)
    assert model_input[0] is x

    checkOut(out)

    with torch.cuda.amp.autocast():
483
        out = model(model_input)
484

Anirudh's avatar
Anirudh committed
485
    checkOut(out)
486

487
488
    _check_input_backprop(model, model_input)

Anirudh's avatar
Anirudh committed
489
490
491
492
    # now switch to cpu and make sure it works
    model.cpu()
    x = x.cpu()
    out_cpu = model([x])
493

Anirudh's avatar
Anirudh committed
494
    checkOut(out_cpu)
495

496
497
    _check_input_backprop(model, [x])

498

Anirudh's avatar
Anirudh committed
499
def test_generalizedrcnn_transform_repr():
500

Anirudh's avatar
Anirudh committed
501
502
503
    min_size, max_size = 224, 299
    image_mean = [0.485, 0.456, 0.406]
    image_std = [0.229, 0.224, 0.225]
504

505
506
507
    t = models.detection.transform.GeneralizedRCNNTransform(
        min_size=min_size, max_size=max_size, image_mean=image_mean, image_std=image_std
    )
508

Anirudh's avatar
Anirudh committed
509
    # Check integrity of object __repr__ attribute
510
511
    expected_string = "GeneralizedRCNNTransform("
    _indent = "\n    "
512
513
    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
514
515
    expected_string += "mode='bilinear')\n)"
    assert t.__repr__() == expected_string
516
517


518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
test_vit_conv_stem_configs = [
    models.vision_transformer.ConvStemConfig(kernel_size=3, stride=2, out_channels=64),
    models.vision_transformer.ConvStemConfig(kernel_size=3, stride=2, out_channels=128),
    models.vision_transformer.ConvStemConfig(kernel_size=3, stride=1, out_channels=128),
    models.vision_transformer.ConvStemConfig(kernel_size=3, stride=2, out_channels=256),
    models.vision_transformer.ConvStemConfig(kernel_size=3, stride=1, out_channels=256),
    models.vision_transformer.ConvStemConfig(kernel_size=3, stride=2, out_channels=512),
]


def vitc_b_16(**kwargs: Any):
    return models.VisionTransformer(
        image_size=224,
        patch_size=16,
        num_layers=12,
        num_heads=12,
        hidden_dim=768,
        mlp_dim=3072,
        conv_stem_configs=test_vit_conv_stem_configs,
        **kwargs,
    )


@pytest.mark.parametrize("model_fn", [vitc_b_16])
@pytest.mark.parametrize("dev", cpu_and_gpu())
def test_vitc_models(model_fn, dev):
    test_classification_model(model_fn, dev)


547
@pytest.mark.parametrize("model_fn", get_models_from_module(models))
548
@pytest.mark.parametrize("dev", cpu_and_gpu())
549
def test_classification_model(model_fn, dev):
Anirudh's avatar
Anirudh committed
550
551
    set_rng_seed(0)
    defaults = {
552
553
        "num_classes": 50,
        "input_shape": (1, 3, 224, 224),
Anirudh's avatar
Anirudh committed
554
    }
555
    model_name = model_fn.__name__
Anirudh's avatar
Anirudh committed
556
    kwargs = {**defaults, **_model_params.get(model_name, {})}
557
    num_classes = kwargs.get("num_classes")
558
    input_shape = kwargs.pop("input_shape")
Anirudh's avatar
Anirudh committed
559

560
    model = model_fn(**kwargs)
Anirudh's avatar
Anirudh committed
561
562
563
564
565
    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)
566
    assert out.shape[-1] == num_classes
Anirudh's avatar
Anirudh committed
567
    _check_jit_scriptable(model, (x,), unwrapper=script_model_unwrapper.get(model_name, None))
568
    _check_fx_compatible(model, x)
Anirudh's avatar
Anirudh committed
569
570
571
572
573
574
575
576

    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
577

578
579
    _check_input_backprop(model, x)

580

581
@pytest.mark.parametrize("model_fn", get_models_from_module(models.segmentation))
582
@pytest.mark.parametrize("dev", cpu_and_gpu())
583
def test_segmentation_model(model_fn, dev):
Anirudh's avatar
Anirudh committed
584
585
    set_rng_seed(0)
    defaults = {
586
587
588
        "num_classes": 10,
        "pretrained_backbone": False,
        "input_shape": (1, 3, 32, 32),
Anirudh's avatar
Anirudh committed
589
    }
590
    model_name = model_fn.__name__
Anirudh's avatar
Anirudh committed
591
    kwargs = {**defaults, **_model_params.get(model_name, {})}
592
    input_shape = kwargs.pop("input_shape")
Anirudh's avatar
Anirudh committed
593

594
    model = model_fn(**kwargs)
Anirudh's avatar
Anirudh committed
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
    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))
621
    _check_fx_compatible(model, x)
Anirudh's avatar
Anirudh committed
622
623
624
625
626
627
628
629
630

    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:
631
        msg = (
632
            f"The output of {test_segmentation_model.__name__} could only be partially validated. "
633
634
            "This is likely due to unit-test flakiness, but you may "
            "want to do additional manual checks if you made "
635
            "significant changes to the codebase."
636
        )
Anirudh's avatar
Anirudh committed
637
638
        warnings.warn(msg, RuntimeWarning)
        pytest.skip(msg)
639

640
641
    _check_input_backprop(model, x)

642

643
@pytest.mark.parametrize("model_fn", get_models_from_module(models.detection))
644
@pytest.mark.parametrize("dev", cpu_and_gpu())
645
def test_detection_model(model_fn, dev):
Anirudh's avatar
Anirudh committed
646
647
    set_rng_seed(0)
    defaults = {
648
649
650
        "num_classes": 50,
        "pretrained_backbone": False,
        "input_shape": (3, 300, 300),
Anirudh's avatar
Anirudh committed
651
    }
652
    model_name = model_fn.__name__
Anirudh's avatar
Anirudh committed
653
    kwargs = {**defaults, **_model_params.get(model_name, {})}
654
    input_shape = kwargs.pop("input_shape")
Anirudh's avatar
Anirudh committed
655

656
    model = model_fn(**kwargs)
Anirudh's avatar
Anirudh committed
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
    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
682
            return tensor[ith_index - 1 :: ith_index]
Anirudh's avatar
Anirudh committed
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704

        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)
705
706
707
            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
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727

            # 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:
728
        msg = (
729
            f"The output of {test_detection_model.__name__} could only be partially validated. "
730
731
            "This is likely due to unit-test flakiness, but you may "
            "want to do additional manual checks if you made "
732
            "significant changes to the codebase."
733
        )
Anirudh's avatar
Anirudh committed
734
735
        warnings.warn(msg, RuntimeWarning)
        pytest.skip(msg)
736

737
738
    _check_input_backprop(model, model_input)

739

740
741
@pytest.mark.parametrize("model_fn", get_models_from_module(models.detection))
def test_detection_model_validation(model_fn):
Anirudh's avatar
Anirudh committed
742
    set_rng_seed(0)
743
    model = model_fn(num_classes=50, pretrained_backbone=False)
Anirudh's avatar
Anirudh committed
744
745
746
747
748
749
750
751
    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
752
    targets = [{"boxes": 0.0}]
Anirudh's avatar
Anirudh committed
753
754
755
756
757
    with pytest.raises(ValueError):
        model(x, targets=targets)

    # validate boxes shape
    for boxes in (torch.rand((4,)), torch.rand((1, 5))):
758
        targets = [{"boxes": boxes}]
Anirudh's avatar
Anirudh committed
759
760
761
762
763
        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]])
764
    targets = [{"boxes": boxes}]
Anirudh's avatar
Anirudh committed
765
766
    with pytest.raises(ValueError):
        model(x, targets=targets)
767

768

769
@pytest.mark.parametrize("model_fn", get_models_from_module(models.video))
770
@pytest.mark.parametrize("dev", cpu_and_gpu())
771
def test_video_model(model_fn, dev):
Anirudh's avatar
Anirudh committed
772
773
774
    # the default input shape is
    # bs * num_channels * clip_len * h *w
    input_shape = (1, 3, 4, 112, 112)
775
    model_name = model_fn.__name__
Anirudh's avatar
Anirudh committed
776
    # test both basicblock and Bottleneck
777
    model = model_fn(num_classes=50)
Anirudh's avatar
Anirudh committed
778
779
780
781
782
    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))
783
    _check_fx_compatible(model, x)
Anirudh's avatar
Anirudh committed
784
785
786
787
788
789
    assert out.shape[-1] == 50

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

791
792
    _check_input_backprop(model, x)

793

794
795
796
797
798
799
800
@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",
)
801
802
@pytest.mark.parametrize("model_fn", get_models_from_module(models.quantization))
def test_quantized_classification_model(model_fn):
803
    set_rng_seed(0)
804
    defaults = {
805
        "num_classes": 5,
806
807
808
        "input_shape": (1, 3, 224, 224),
        "pretrained": False,
        "quantize": True,
809
    }
810
    model_name = model_fn.__name__
811
    kwargs = {**defaults, **_model_params.get(model_name, {})}
812
    input_shape = kwargs.pop("input_shape")
813
814

    # First check if quantize=True provides models that can run with input data
815
    model = model_fn(**kwargs)
816
    model.eval()
817
    x = torch.rand(input_shape)
818
819
820
821
822
823
824
    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)
825

826
    kwargs["quantize"] = False
827
    for eval_mode in [True, False]:
828
        model = model_fn(**kwargs)
829
830
        if eval_mode:
            model.eval()
831
            model.qconfig = torch.ao.quantization.default_qconfig
832
833
        else:
            model.train()
834
            model.qconfig = torch.ao.quantization.default_qat_qconfig
835

836
        model.fuse_model(is_qat=not eval_mode)
837
        if eval_mode:
838
            torch.ao.quantization.prepare(model, inplace=True)
839
        else:
840
            torch.ao.quantization.prepare_qat(model, inplace=True)
841
842
            model.eval()

843
        torch.ao.quantization.convert(model, inplace=True)
844
845
846
847
848
849
850
851

    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


852
@pytest.mark.parametrize("model_fn", get_models_from_module(models.detection))
853
def test_detection_model_trainable_backbone_layers(model_fn, disable_weight_loading):
854
    model_name = model_fn.__name__
855
856
857
    max_trainable = _model_tests_values[model_name]["max_trainable"]
    n_trainable_params = []
    for trainable_layers in range(0, max_trainable + 1):
858
        model = model_fn(pretrained=False, pretrained_backbone=True, trainable_backbone_layers=trainable_layers)
859
860
861
862
863

        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"]


864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
@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)


892
if __name__ == "__main__":
893
    pytest.main([__file__])