test_models.py 31.2 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, 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
133

    sm = torch.jit.script(nn_module)

134
135
    if eager_out is None:
        with torch.no_grad(), freeze_rng_state():
136
            eager_out = nn_module(*args)
137

138
    with torch.no_grad(), freeze_rng_state():
139
140
141
142
143
        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)
144
145
146
147
148
149
150
151

    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)
152
153


154
def _check_fx_compatible(model, inputs, eager_out=None):
155
    model_fx = torch.fx.symbolic_trace(model)
156
157
158
159
    if eager_out is None:
        eager_out = model(inputs)
    fx_out = model_fx(inputs)
    torch.testing.assert_close(eager_out, fx_out)
160
161


162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
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)


191
192
193
# 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
194
script_model_unwrapper = {
195
196
    "googlenet": lambda x: x.logits,
    "inception_v3": lambda x: x.logits,
197
    "fasterrcnn_resnet50_fpn": lambda x: x[1],
198
    "fasterrcnn_mobilenet_v3_large_fpn": lambda x: x[1],
199
    "fasterrcnn_mobilenet_v3_large_320_fpn": lambda x: x[1],
200
201
202
    "maskrcnn_resnet50_fpn": lambda x: x[1],
    "keypointrcnn_resnet50_fpn": lambda x: x[1],
    "retinanet_resnet50_fpn": lambda x: x[1],
203
    "ssd300_vgg16": lambda x: x[1],
204
    "ssdlite320_mobilenet_v3_large": lambda x: x[1],
Hu Ye's avatar
Hu Ye committed
205
    "fcos_resnet50_fpn": lambda x: x[1],
206
}
207
208


209
210
211
212
213
214
215
216
217
218
219
220
221
222
# 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",
223
224
    "deeplabv3_resnet50",
    "deeplabv3_resnet101",
225
    "deeplabv3_mobilenet_v3_large",
226
227
    "fcn_resnet50",
    "fcn_resnet101",
228
    "lraspp_mobilenet_v3_large",
229
    "maskrcnn_resnet50_fpn",
230
231
)

232
233
234
# 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.
235
quantized_flaky_models = ("inception_v3", "resnet50")
236

237

238
239
240
# The following contains configuration parameters for all models which are used by
# the _test_*_model methods.
_model_params = {
241
242
243
244
245
246
247
    "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),
248
    },
249
250
251
252
253
254
    "keypointrcnn_resnet50_fpn": {
        "num_classes": 2,
        "min_size": 224,
        "max_size": 224,
        "box_score_thresh": 0.15,
        "input_shape": (3, 224, 224),
255
    },
256
257
258
259
260
    "fasterrcnn_resnet50_fpn": {
        "num_classes": 20,
        "min_size": 224,
        "max_size": 224,
        "input_shape": (3, 224, 224),
261
    },
Hu Ye's avatar
Hu Ye committed
262
263
264
265
266
267
268
    "fcos_resnet50_fpn": {
        "num_classes": 2,
        "score_thresh": 0.05,
        "min_size": 224,
        "max_size": 224,
        "input_shape": (3, 224, 224),
    },
269
270
271
272
273
    "maskrcnn_resnet50_fpn": {
        "num_classes": 10,
        "min_size": 224,
        "max_size": 224,
        "input_shape": (3, 224, 224),
274
    },
275
276
    "fasterrcnn_mobilenet_v3_large_fpn": {
        "box_score_thresh": 0.02076,
277
    },
278
279
280
281
    "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,
282
    },
283
284
285
286
    "vit_h_14": {
        "image_size": 56,
        "input_shape": (1, 3, 56, 56),
    },
287
}
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
# 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)}
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
337
338
339
340
341
# 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
342
343
344
345
    "fcos_resnet50_fpn": {
        "max_trainable": 5,
        "n_trn_params_per_layer": [54, 64, 83, 96, 106, 107],
    },
346
347
348
}


Anirudh's avatar
Anirudh committed
349
350
351
352
353
354
355
356
357
358
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


359
360
@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
361
362
363
    input_shape = (1, 3, 300, 300)
    x = torch.rand(input_shape)

364
    model1 = model_fn(num_classes=50, memory_efficient=True)
Anirudh's avatar
Anirudh committed
365
    params = model1.state_dict()
366
    num_params = sum(x.numel() for x in model1.parameters())
Anirudh's avatar
Anirudh committed
367
368
369
    model1.eval()
    out1 = model1(x)
    out1.sum().backward()
370
    num_grad = sum(x.grad.numel() for x in model1.parameters() if x.grad is not None)
Anirudh's avatar
Anirudh committed
371

372
    model2 = model_fn(num_classes=50, memory_efficient=False)
Anirudh's avatar
Anirudh committed
373
374
375
376
377
378
379
    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)

380
381
382
    _check_input_backprop(model1, x)
    _check_input_backprop(model2, x)

Anirudh's avatar
Anirudh committed
383

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


406
407
408
@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
409
410
411
    assert any(isinstance(x, nn.BatchNorm2d) for x in model.modules())

    def get_gn(num_channels):
412
        return nn.GroupNorm(1, num_channels)
Anirudh's avatar
Anirudh committed
413

414
    model = model_fn(norm_layer=get_gn)
415
    assert not (any(isinstance(x, nn.BatchNorm2d) for x in model.modules()))
Anirudh's avatar
Anirudh committed
416
417
418
419
420
    assert any(isinstance(x, nn.GroupNorm) for x in model.modules())


def test_inception_v3_eval():
    kwargs = {}
421
422
423
    kwargs["transform_input"] = True
    kwargs["aux_logits"] = True
    kwargs["init_weights"] = False
Anirudh's avatar
Anirudh committed
424
425
426
427
428
429
430
    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))
431
    _check_input_backprop(model, x)
Anirudh's avatar
Anirudh committed
432
433
434


def test_fasterrcnn_double():
435
    model = models.detection.fasterrcnn_resnet50_fpn(num_classes=50, weights=None, weights_backbone=None)
Anirudh's avatar
Anirudh committed
436
437
438
439
440
441
442
443
444
445
446
    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]
447
    _check_input_backprop(model, model_input)
Anirudh's avatar
Anirudh committed
448
449
450
451


def test_googlenet_eval():
    kwargs = {}
452
453
454
    kwargs["transform_input"] = True
    kwargs["aux_logits"] = True
    kwargs["init_weights"] = False
Anirudh's avatar
Anirudh committed
455
456
457
458
459
460
461
462
    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))
463
    _check_input_backprop(model, x)
Anirudh's avatar
Anirudh committed
464
465
466
467
468
469
470
471
472
473


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

474
    model = models.detection.fasterrcnn_resnet50_fpn(num_classes=50, weights=None, weights_backbone=None)
Anirudh's avatar
Anirudh committed
475
476
477
    model.cuda()
    model.eval()
    input_shape = (3, 300, 300)
478
    x = torch.rand(input_shape, device="cuda")
Anirudh's avatar
Anirudh committed
479
480
481
482
483
484
485
    model_input = [x]
    out = model(model_input)
    assert model_input[0] is x

    checkOut(out)

    with torch.cuda.amp.autocast():
486
        out = model(model_input)
487

Anirudh's avatar
Anirudh committed
488
    checkOut(out)
489

490
491
    _check_input_backprop(model, model_input)

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

Anirudh's avatar
Anirudh committed
497
    checkOut(out_cpu)
498

499
500
    _check_input_backprop(model, [x])

501

Anirudh's avatar
Anirudh committed
502
def test_generalizedrcnn_transform_repr():
503

Anirudh's avatar
Anirudh committed
504
505
506
    min_size, max_size = 224, 299
    image_mean = [0.485, 0.456, 0.406]
    image_std = [0.229, 0.224, 0.225]
507

508
509
510
    t = models.detection.transform.GeneralizedRCNNTransform(
        min_size=min_size, max_size=max_size, image_mean=image_mean, image_std=image_std
    )
511

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


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

563
    model = model_fn(**kwargs)
Anirudh's avatar
Anirudh committed
564
565
566
567
568
    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)
569
    assert out.shape[-1] == num_classes
570
571
    _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
572
573
574
575
576
577
578
579

    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
580

581
582
    _check_input_backprop(model, x)

583

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

597
    model = model_fn(**kwargs)
Anirudh's avatar
Anirudh committed
598
599
600
    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)
601
    out = model(x)
Anirudh's avatar
Anirudh committed
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620

    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

621
    full_validation = check_out(out["out"])
Anirudh's avatar
Anirudh committed
622

623
624
    _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
625
626
627

    if dev == torch.device("cuda"):
        with torch.cuda.amp.autocast():
628
            out = model(x)
Anirudh's avatar
Anirudh committed
629
630
            # See autocast_flaky_numerics comment at top of file.
            if model_name not in autocast_flaky_numerics:
631
                full_validation &= check_out(out["out"])
Anirudh's avatar
Anirudh committed
632
633

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

643
644
    _check_input_backprop(model, x)

645

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

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

        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)
709
710
711
            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
712
713
714
715
716
717
718
719
720
721

            # 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)
722
    _check_jit_scriptable(model, ([x],), unwrapper=script_model_unwrapper.get(model_name, None), eager_out=out)
Anirudh's avatar
Anirudh committed
723
724
725
726
727
728
729
730
731

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

741
742
    _check_input_backprop(model, model_input)

743

744
745
@pytest.mark.parametrize("model_fn", get_models_from_module(models.detection))
def test_detection_model_validation(model_fn):
Anirudh's avatar
Anirudh committed
746
    set_rng_seed(0)
747
    model = model_fn(num_classes=50, weights=None, weights_backbone=None)
Anirudh's avatar
Anirudh committed
748
749
750
751
    input_shape = (3, 300, 300)
    x = [torch.rand(input_shape)]

    # validate that targets are present in training
752
    with pytest.raises(AssertionError):
Anirudh's avatar
Anirudh committed
753
754
755
        model(x)

    # validate type
756
    targets = [{"boxes": 0.0}]
757
    with pytest.raises(AssertionError):
Anirudh's avatar
Anirudh committed
758
759
760
761
        model(x, targets=targets)

    # validate boxes shape
    for boxes in (torch.rand((4,)), torch.rand((1, 5))):
762
        targets = [{"boxes": boxes}]
763
        with pytest.raises(AssertionError):
Anirudh's avatar
Anirudh committed
764
765
766
767
            model(x, targets=targets)

    # validate that no degenerate boxes are present
    boxes = torch.tensor([[1, 3, 1, 4], [2, 4, 3, 4]])
768
    targets = [{"boxes": boxes}]
769
    with pytest.raises(AssertionError):
Anirudh's avatar
Anirudh committed
770
        model(x, targets=targets)
771

772

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

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

795
796
    _check_input_backprop(model, x)

797

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

    # First check if quantize=True provides models that can run with input data
818
    model = model_fn(**kwargs)
819
    model.eval()
820
    x = torch.rand(input_shape)
821
822
823
824
825
    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
826
827
828
829
830
831
832
        _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
833

834
    kwargs["quantize"] = False
835
    for eval_mode in [True, False]:
836
        model = model_fn(**kwargs)
837
838
        if eval_mode:
            model.eval()
839
            model.qconfig = torch.ao.quantization.default_qconfig
840
841
        else:
            model.train()
842
            model.qconfig = torch.ao.quantization.default_qat_qconfig
843

844
        model.fuse_model(is_qat=not eval_mode)
845
        if eval_mode:
846
            torch.ao.quantization.prepare(model, inplace=True)
847
        else:
848
            torch.ao.quantization.prepare_qat(model, inplace=True)
849
850
            model.eval()

851
        torch.ao.quantization.convert(model, inplace=True)
852
853


854
@pytest.mark.parametrize("model_fn", get_models_from_module(models.detection))
855
def test_detection_model_trainable_backbone_layers(model_fn, disable_weight_loading):
856
    model_name = model_fn.__name__
857
858
859
    max_trainable = _model_tests_values[model_name]["max_trainable"]
    n_trainable_params = []
    for trainable_layers in range(0, max_trainable + 1):
860
        model = model_fn(weights=None, weights_backbone="DEFAULT", trainable_backbone_layers=trainable_layers)
861
862
863
864
865

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


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
892
893
@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)


894
if __name__ == "__main__":
895
    pytest.main([__file__])