test_models.py 32 KB
Newer Older
1
import contextlib
2
3
4
import functools
import operator
import os
5
6
import pkgutil
import sys
7
import warnings
8
from collections import OrderedDict
9
from tempfile import TemporaryDirectory
10
from typing import Any
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 = 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
        torch.testing.assert_close(output, expected, rtol=rtol, atol=atol, check_dtype=False)


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

124
125
126
127
128
129
130
    def get_export_import_copy(m):
        """Save and load a TorchScript model"""
        with TemporaryDirectory() as dir:
            path = os.path.join(dir, "script.pt")
            m.save(path)
            imported = torch.jit.load(path)
        return imported
131

132
    TEST_WITH_SLOW = os.getenv("PYTORCH_TEST_WITH_SLOW", "0") == "1"
133
134
    if not TEST_WITH_SLOW or skip:
        # TorchScript is not enabled, skip these tests
135
        msg = (
136
            f"The check_jit_scriptable test for {nn_module.__class__.__name__} was skipped. "
137
138
139
140
            "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 "
141
            "manually skipped."
142
        )
143
144
145
146
147
        warnings.warn(msg, RuntimeWarning)
        return None

    sm = torch.jit.script(nn_module)

148
149
150
151
    if eager_out is None:
        with torch.no_grad(), freeze_rng_state():
            if unwrapper:
                eager_out = nn_module(*args)
152

153
    with torch.no_grad(), freeze_rng_state():
154
155
156
157
158
        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)
159
160
161
162
163
164
165
166

    m_import = get_export_import_copy(sm)
    with torch.no_grad(), freeze_rng_state():
        imported_script_out = m_import(*args)
        if unwrapper:
            imported_script_out = unwrapper(imported_script_out)

    torch.testing.assert_close(script_out, imported_script_out, atol=3e-4, rtol=3e-4)
167
168


169
def _check_fx_compatible(model, inputs, eager_out=None):
170
    model_fx = torch.fx.symbolic_trace(model)
171
172
173
174
    if eager_out is None:
        eager_out = model(inputs)
    fx_out = model_fx(inputs)
    torch.testing.assert_close(eager_out, fx_out)
175
176


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


206
207
208
# 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
209
script_model_unwrapper = {
210
211
    "googlenet": lambda x: x.logits,
    "inception_v3": lambda x: x.logits,
212
    "fasterrcnn_resnet50_fpn": lambda x: x[1],
213
    "fasterrcnn_mobilenet_v3_large_fpn": lambda x: x[1],
214
    "fasterrcnn_mobilenet_v3_large_320_fpn": lambda x: x[1],
215
216
217
    "maskrcnn_resnet50_fpn": lambda x: x[1],
    "keypointrcnn_resnet50_fpn": lambda x: x[1],
    "retinanet_resnet50_fpn": lambda x: x[1],
218
    "ssd300_vgg16": lambda x: x[1],
219
    "ssdlite320_mobilenet_v3_large": lambda x: x[1],
Hu Ye's avatar
Hu Ye committed
220
    "fcos_resnet50_fpn": 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
    },
Hu Ye's avatar
Hu Ye committed
277
278
279
280
281
282
283
    "fcos_resnet50_fpn": {
        "num_classes": 2,
        "score_thresh": 0.05,
        "min_size": 224,
        "max_size": 224,
        "input_shape": (3, 224, 224),
    },
284
285
286
287
288
    "maskrcnn_resnet50_fpn": {
        "num_classes": 10,
        "min_size": 224,
        "max_size": 224,
        "input_shape": (3, 224, 224),
289
    },
290
291
    "fasterrcnn_mobilenet_v3_large_fpn": {
        "box_score_thresh": 0.02076,
292
    },
293
294
295
296
    "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,
297
298
    },
}
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
# speeding up slow models:
slow_models = [
    "convnext_base",
    "convnext_large",
    "resnext101_32x8d",
    "wide_resnet101_2",
    "efficientnet_b6",
    "efficientnet_b7",
    "efficientnet_v2_m",
    "efficientnet_v2_l",
    "regnet_y_16gf",
    "regnet_y_32gf",
    "regnet_y_128gf",
    "regnet_x_16gf",
    "regnet_x_32gf",
]
for m in slow_models:
    _model_params[m] = {"input_shape": (1, 3, 64, 64)}
317
318


319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
# 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
353
354
355
356
    "fcos_resnet50_fpn": {
        "max_trainable": 5,
        "n_trn_params_per_layer": [54, 64, 83, 96, 106, 107],
    },
357
358
359
}


Anirudh's avatar
Anirudh committed
360
361
362
363
364
365
366
367
368
369
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


370
371
@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
372
373
374
    input_shape = (1, 3, 300, 300)
    x = torch.rand(input_shape)

375
    model1 = model_fn(num_classes=50, memory_efficient=True)
Anirudh's avatar
Anirudh committed
376
    params = model1.state_dict()
377
    num_params = sum(x.numel() for x in model1.parameters())
Anirudh's avatar
Anirudh committed
378
379
380
    model1.eval()
    out1 = model1(x)
    out1.sum().backward()
381
    num_grad = sum(x.grad.numel() for x in model1.parameters() if x.grad is not None)
Anirudh's avatar
Anirudh committed
382

383
    model2 = model_fn(num_classes=50, memory_efficient=False)
Anirudh's avatar
Anirudh committed
384
385
386
387
388
389
390
    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)

391
392
393
    _check_input_backprop(model1, x)
    _check_input_backprop(model2, x)

Anirudh's avatar
Anirudh committed
394

395
396
397
@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
398
399
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
400
    model = models.resnet50(replace_stride_with_dilation=(dilate_layer_2, dilate_layer_3, dilate_layer_4))
Anirudh's avatar
Anirudh committed
401
402
403
404
405
406
407
408
409
    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():
410
    model = models.mobilenet_v2(inverted_residual_setting=[[1, 16, 1, 1], [6, 24, 2, 2]])
Anirudh's avatar
Anirudh committed
411
412
413
414
415
416
    model.eval()
    x = torch.rand(1, 3, 224, 224)
    out = model(x)
    assert out.shape[-1] == 1000


417
418
419
@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
420
421
422
423
424
    assert any(isinstance(x, nn.BatchNorm2d) for x in model.modules())

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

425
    model = model_fn(norm_layer=get_gn)
426
    assert not (any(isinstance(x, nn.BatchNorm2d) for x in model.modules()))
Anirudh's avatar
Anirudh committed
427
428
429
430
431
432
    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 = {}
433
434
435
    kwargs["transform_input"] = True
    kwargs["aux_logits"] = True
    kwargs["init_weights"] = False
Anirudh's avatar
Anirudh committed
436
437
438
439
440
441
442
    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))
443
    _check_input_backprop(model, x)
Anirudh's avatar
Anirudh committed
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458


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]
459
    _check_input_backprop(model, model_input)
Anirudh's avatar
Anirudh committed
460
461
462
463
464


def test_googlenet_eval():
    # replacement for models.googlenet(pretrained=True) that does not download weights
    kwargs = {}
465
466
467
    kwargs["transform_input"] = True
    kwargs["aux_logits"] = True
    kwargs["init_weights"] = False
Anirudh's avatar
Anirudh committed
468
469
470
471
472
473
474
475
    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))
476
    _check_input_backprop(model, x)
Anirudh's avatar
Anirudh committed
477
478
479
480
481
482
483
484
485
486
487
488
489
490


@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)
491
    x = torch.rand(input_shape, device="cuda")
Anirudh's avatar
Anirudh committed
492
493
494
495
496
497
498
    model_input = [x]
    out = model(model_input)
    assert model_input[0] is x

    checkOut(out)

    with torch.cuda.amp.autocast():
499
        out = model(model_input)
500

Anirudh's avatar
Anirudh committed
501
    checkOut(out)
502

503
504
    _check_input_backprop(model, model_input)

Anirudh's avatar
Anirudh committed
505
506
507
508
    # now switch to cpu and make sure it works
    model.cpu()
    x = x.cpu()
    out_cpu = model([x])
509

Anirudh's avatar
Anirudh committed
510
    checkOut(out_cpu)
511

512
513
    _check_input_backprop(model, [x])

514

Anirudh's avatar
Anirudh committed
515
def test_generalizedrcnn_transform_repr():
516

Anirudh's avatar
Anirudh committed
517
518
519
    min_size, max_size = 224, 299
    image_mean = [0.485, 0.456, 0.406]
    image_std = [0.229, 0.224, 0.225]
520

521
522
523
    t = models.detection.transform.GeneralizedRCNNTransform(
        min_size=min_size, max_size=max_size, image_mean=image_mean, image_std=image_std
    )
524

Anirudh's avatar
Anirudh committed
525
    # Check integrity of object __repr__ attribute
526
527
    expected_string = "GeneralizedRCNNTransform("
    _indent = "\n    "
528
529
    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
530
531
    expected_string += "mode='bilinear')\n)"
    assert t.__repr__() == expected_string
532
533


534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
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)


563
@pytest.mark.parametrize("model_fn", get_models_from_module(models))
564
@pytest.mark.parametrize("dev", cpu_and_gpu())
565
def test_classification_model(model_fn, dev):
Anirudh's avatar
Anirudh committed
566
567
    set_rng_seed(0)
    defaults = {
568
569
        "num_classes": 50,
        "input_shape": (1, 3, 224, 224),
Anirudh's avatar
Anirudh committed
570
    }
571
    model_name = model_fn.__name__
Anirudh's avatar
Anirudh committed
572
    kwargs = {**defaults, **_model_params.get(model_name, {})}
573
    num_classes = kwargs.get("num_classes")
574
    input_shape = kwargs.pop("input_shape")
Anirudh's avatar
Anirudh committed
575

576
    model = model_fn(**kwargs)
Anirudh's avatar
Anirudh committed
577
578
579
580
581
    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)
582
    assert out.shape[-1] == num_classes
583
584
    _check_jit_scriptable(model, (x,), unwrapper=script_model_unwrapper.get(model_name, None), eager_out=out)
    _check_fx_compatible(model, x, eager_out=out)
Anirudh's avatar
Anirudh committed
585
586
587
588
589
590
591
592

    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
593

594
595
    _check_input_backprop(model, x)

596

597
@pytest.mark.parametrize("model_fn", get_models_from_module(models.segmentation))
598
@pytest.mark.parametrize("dev", cpu_and_gpu())
599
def test_segmentation_model(model_fn, dev):
Anirudh's avatar
Anirudh committed
600
601
    set_rng_seed(0)
    defaults = {
602
603
604
        "num_classes": 10,
        "pretrained_backbone": False,
        "input_shape": (1, 3, 32, 32),
Anirudh's avatar
Anirudh committed
605
    }
606
    model_name = model_fn.__name__
Anirudh's avatar
Anirudh committed
607
    kwargs = {**defaults, **_model_params.get(model_name, {})}
608
    input_shape = kwargs.pop("input_shape")
Anirudh's avatar
Anirudh committed
609

610
    model = model_fn(**kwargs)
Anirudh's avatar
Anirudh committed
611
612
613
    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)
614
    out = model(x)
Anirudh's avatar
Anirudh committed
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633

    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

634
    full_validation = check_out(out["out"])
Anirudh's avatar
Anirudh committed
635

636
637
    _check_jit_scriptable(model, (x,), unwrapper=script_model_unwrapper.get(model_name, None), eager_out=out)
    _check_fx_compatible(model, x, eager_out=out)
Anirudh's avatar
Anirudh committed
638
639
640

    if dev == torch.device("cuda"):
        with torch.cuda.amp.autocast():
641
            out = model(x)
Anirudh's avatar
Anirudh committed
642
643
            # See autocast_flaky_numerics comment at top of file.
            if model_name not in autocast_flaky_numerics:
644
                full_validation &= check_out(out["out"])
Anirudh's avatar
Anirudh committed
645
646

    if not full_validation:
647
        msg = (
648
            f"The output of {test_segmentation_model.__name__} could only be partially validated. "
649
650
            "This is likely due to unit-test flakiness, but you may "
            "want to do additional manual checks if you made "
651
            "significant changes to the codebase."
652
        )
Anirudh's avatar
Anirudh committed
653
654
        warnings.warn(msg, RuntimeWarning)
        pytest.skip(msg)
655

656
657
    _check_input_backprop(model, x)

658

659
@pytest.mark.parametrize("model_fn", get_models_from_module(models.detection))
660
@pytest.mark.parametrize("dev", cpu_and_gpu())
661
def test_detection_model(model_fn, dev):
Anirudh's avatar
Anirudh committed
662
663
    set_rng_seed(0)
    defaults = {
664
665
666
        "num_classes": 50,
        "pretrained_backbone": False,
        "input_shape": (3, 300, 300),
Anirudh's avatar
Anirudh committed
667
    }
668
    model_name = model_fn.__name__
Anirudh's avatar
Anirudh committed
669
    kwargs = {**defaults, **_model_params.get(model_name, {})}
670
    input_shape = kwargs.pop("input_shape")
Anirudh's avatar
Anirudh committed
671

672
    model = model_fn(**kwargs)
Anirudh's avatar
Anirudh committed
673
674
675
676
677
678
679
680
681
682
683
    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):
684
            tensor = tensor.cpu()
Anirudh's avatar
Anirudh committed
685
686
687
688
689
690
691
692
693
694
695
696
697
698
            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
699
            return tensor[ith_index - 1 :: ith_index]
Anirudh's avatar
Anirudh committed
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721

        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)
722
723
724
            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
725
726
727
728
729
730
731
732
733
734

            # 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)
735
    _check_jit_scriptable(model, ([x],), unwrapper=script_model_unwrapper.get(model_name, None), eager_out=out)
Anirudh's avatar
Anirudh committed
736
737
738
739
740
741
742
743
744

    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:
745
        msg = (
746
            f"The output of {test_detection_model.__name__} could only be partially validated. "
747
748
            "This is likely due to unit-test flakiness, but you may "
            "want to do additional manual checks if you made "
749
            "significant changes to the codebase."
750
        )
Anirudh's avatar
Anirudh committed
751
752
        warnings.warn(msg, RuntimeWarning)
        pytest.skip(msg)
753

754
755
    _check_input_backprop(model, model_input)

756

757
758
@pytest.mark.parametrize("model_fn", get_models_from_module(models.detection))
def test_detection_model_validation(model_fn):
Anirudh's avatar
Anirudh committed
759
    set_rng_seed(0)
760
    model = model_fn(num_classes=50, pretrained_backbone=False)
Anirudh's avatar
Anirudh committed
761
762
763
764
765
766
767
768
    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
769
    targets = [{"boxes": 0.0}]
770
    with pytest.raises(TypeError):
Anirudh's avatar
Anirudh committed
771
772
773
774
        model(x, targets=targets)

    # validate boxes shape
    for boxes in (torch.rand((4,)), torch.rand((1, 5))):
775
        targets = [{"boxes": boxes}]
Anirudh's avatar
Anirudh committed
776
777
778
779
780
        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]])
781
    targets = [{"boxes": boxes}]
Anirudh's avatar
Anirudh committed
782
783
    with pytest.raises(ValueError):
        model(x, targets=targets)
784

785

786
@pytest.mark.parametrize("model_fn", get_models_from_module(models.video))
787
@pytest.mark.parametrize("dev", cpu_and_gpu())
788
def test_video_model(model_fn, dev):
Anirudh's avatar
Anirudh committed
789
790
791
    # the default input shape is
    # bs * num_channels * clip_len * h *w
    input_shape = (1, 3, 4, 112, 112)
792
    model_name = model_fn.__name__
Anirudh's avatar
Anirudh committed
793
    # test both basicblock and Bottleneck
794
    model = model_fn(num_classes=50)
Anirudh's avatar
Anirudh committed
795
796
797
798
    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)
799
800
    _check_jit_scriptable(model, (x,), unwrapper=script_model_unwrapper.get(model_name, None), eager_out=out)
    _check_fx_compatible(model, x, eager_out=out)
Anirudh's avatar
Anirudh committed
801
802
803
804
805
806
    assert out.shape[-1] == 50

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

808
809
    _check_input_backprop(model, x)

810

811
812
813
814
815
816
817
@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",
)
818
819
@pytest.mark.parametrize("model_fn", get_models_from_module(models.quantization))
def test_quantized_classification_model(model_fn):
820
    set_rng_seed(0)
821
    defaults = {
822
        "num_classes": 5,
823
824
825
        "input_shape": (1, 3, 224, 224),
        "pretrained": False,
        "quantize": True,
826
    }
827
    model_name = model_fn.__name__
828
    kwargs = {**defaults, **_model_params.get(model_name, {})}
829
    input_shape = kwargs.pop("input_shape")
830
831

    # First check if quantize=True provides models that can run with input data
832
    model = model_fn(**kwargs)
833
    model.eval()
834
    x = torch.rand(input_shape)
835
836
837
838
839
    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
840
841
842
843
844
845
846
        _check_jit_scriptable(model, (x,), unwrapper=script_model_unwrapper.get(model_name, None), eager_out=out)
        _check_fx_compatible(model, x, eager_out=out)
    else:
        try:
            torch.jit.script(model)
        except Exception as e:
            raise AssertionError("model cannot be scripted.") from e
847

848
    kwargs["quantize"] = False
849
    for eval_mode in [True, False]:
850
        model = model_fn(**kwargs)
851
852
        if eval_mode:
            model.eval()
853
            model.qconfig = torch.ao.quantization.default_qconfig
854
855
        else:
            model.train()
856
            model.qconfig = torch.ao.quantization.default_qat_qconfig
857

858
        model.fuse_model(is_qat=not eval_mode)
859
        if eval_mode:
860
            torch.ao.quantization.prepare(model, inplace=True)
861
        else:
862
            torch.ao.quantization.prepare_qat(model, inplace=True)
863
864
            model.eval()

865
        torch.ao.quantization.convert(model, inplace=True)
866
867


868
@pytest.mark.parametrize("model_fn", get_models_from_module(models.detection))
869
def test_detection_model_trainable_backbone_layers(model_fn, disable_weight_loading):
870
    model_name = model_fn.__name__
871
872
873
    max_trainable = _model_tests_values[model_name]["max_trainable"]
    n_trainable_params = []
    for trainable_layers in range(0, max_trainable + 1):
874
        model = model_fn(pretrained=False, pretrained_backbone=True, trainable_backbone_layers=trainable_layers)
875
876
877
878
879

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


880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
@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)


908
if __name__ == "__main__":
909
    pytest.main([__file__])