test_models.py 33.4 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
from _utils_internal import get_relative_path
17
from common_utils import cpu_and_gpu, freeze_rng_state, map_nested_tensor_object, needs_cuda, set_rng_seed
18
from torchvision import models
19

20
ACCEPT = os.getenv("EXPECTTEST_ACCEPT", "0") == "1"
21
SKIP_BIG_MODEL = os.getenv("SKIP_BIG_MODEL", "1") == "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 = 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
        torch.testing.assert_close(output, expected, rtol=rtol, atol=atol, check_dtype=False)


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

125
126
127
128
129
130
131
    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
132
133
134

    sm = torch.jit.script(nn_module)

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

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

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


155
def _check_fx_compatible(model, inputs, eager_out=None):
156
    model_fx = torch.fx.symbolic_trace(model)
157
158
159
160
    if eager_out is None:
        eager_out = model(inputs)
    fx_out = model_fx(inputs)
    torch.testing.assert_close(eager_out, fx_out)
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
191
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)


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


213
214
215
216
217
218
219
220
221
222
223
224
225
226
# 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",
227
228
    "deeplabv3_resnet50",
    "deeplabv3_resnet101",
229
    "deeplabv3_mobilenet_v3_large",
230
231
    "fcn_resnet50",
    "fcn_resnet101",
232
    "lraspp_mobilenet_v3_large",
233
    "maskrcnn_resnet50_fpn",
234
    "maskrcnn_resnet50_fpn_v2",
235
    "keypointrcnn_resnet50_fpn",
236
237
)

238
239
240
# 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.
241
quantized_flaky_models = ("inception_v3", "resnet50")
242

243

244
245
246
# The following contains configuration parameters for all models which are used by
# the _test_*_model methods.
_model_params = {
247
248
249
250
251
252
253
    "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),
254
    },
255
256
257
258
259
260
261
    "retinanet_resnet50_fpn_v2": {
        "num_classes": 20,
        "score_thresh": 0.01,
        "min_size": 224,
        "max_size": 224,
        "input_shape": (3, 224, 224),
    },
262
263
264
265
    "keypointrcnn_resnet50_fpn": {
        "num_classes": 2,
        "min_size": 224,
        "max_size": 224,
266
        "box_score_thresh": 0.17,
267
        "input_shape": (3, 224, 224),
268
    },
269
270
271
272
273
    "fasterrcnn_resnet50_fpn": {
        "num_classes": 20,
        "min_size": 224,
        "max_size": 224,
        "input_shape": (3, 224, 224),
274
    },
275
276
277
278
279
280
    "fasterrcnn_resnet50_fpn_v2": {
        "num_classes": 20,
        "min_size": 224,
        "max_size": 224,
        "input_shape": (3, 224, 224),
    },
Hu Ye's avatar
Hu Ye committed
281
282
283
284
285
286
287
    "fcos_resnet50_fpn": {
        "num_classes": 2,
        "score_thresh": 0.05,
        "min_size": 224,
        "max_size": 224,
        "input_shape": (3, 224, 224),
    },
288
289
290
291
292
    "maskrcnn_resnet50_fpn": {
        "num_classes": 10,
        "min_size": 224,
        "max_size": 224,
        "input_shape": (3, 224, 224),
293
    },
294
295
296
297
298
299
    "maskrcnn_resnet50_fpn_v2": {
        "num_classes": 10,
        "min_size": 224,
        "max_size": 224,
        "input_shape": (3, 224, 224),
    },
300
301
    "fasterrcnn_mobilenet_v3_large_fpn": {
        "box_score_thresh": 0.02076,
302
    },
303
304
305
306
    "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,
307
    },
308
309
310
311
    "vit_h_14": {
        "image_size": 56,
        "input_shape": (1, 3, 56, 56),
    },
312
313
314
    "mvit_v1_b": {
        "input_shape": (1, 3, 16, 224, 224),
    },
315
}
316
317
318
319
320
# speeding up slow models:
slow_models = [
    "convnext_base",
    "convnext_large",
    "resnext101_32x8d",
321
    "resnext101_64x4d",
322
323
324
325
326
327
328
329
330
331
    "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",
Joao Gomes's avatar
Joao Gomes committed
332
    "swin_t",
333
334
    "swin_s",
    "swin_b",
335
336
337
]
for m in slow_models:
    _model_params[m] = {"input_shape": (1, 3, 64, 64)}
338
339


340
341
342
343
344
345
# skip big models to reduce memory usage on CI test
skipped_big_models = {
    "vit_h_14",
    "regnet_y_128gf",
}

346
347
348
349
350
351
# 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],
    },
352
353
354
355
    "retinanet_resnet50_fpn_v2": {
        "max_trainable": 5,
        "n_trn_params_per_layer": [44, 74, 131, 170, 200, 203],
    },
356
357
358
359
360
361
362
363
    "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],
    },
364
365
366
367
    "fasterrcnn_resnet50_fpn_v2": {
        "max_trainable": 5,
        "n_trn_params_per_layer": [50, 80, 137, 176, 206, 209],
    },
368
369
370
371
    "maskrcnn_resnet50_fpn": {
        "max_trainable": 5,
        "n_trn_params_per_layer": [42, 52, 71, 84, 94, 95],
    },
372
373
374
375
    "maskrcnn_resnet50_fpn_v2": {
        "max_trainable": 5,
        "n_trn_params_per_layer": [66, 96, 153, 192, 222, 225],
    },
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
    "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
392
393
394
395
    "fcos_resnet50_fpn": {
        "max_trainable": 5,
        "n_trn_params_per_layer": [54, 64, 83, 96, 106, 107],
    },
396
397
398
}


Anirudh's avatar
Anirudh committed
399
400
401
402
403
404
405
406
407
408
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


409
410
@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
411
412
413
    input_shape = (1, 3, 300, 300)
    x = torch.rand(input_shape)

414
    model1 = model_fn(num_classes=50, memory_efficient=True)
Anirudh's avatar
Anirudh committed
415
    params = model1.state_dict()
416
    num_params = sum(x.numel() for x in model1.parameters())
Anirudh's avatar
Anirudh committed
417
418
419
    model1.eval()
    out1 = model1(x)
    out1.sum().backward()
420
    num_grad = sum(x.grad.numel() for x in model1.parameters() if x.grad is not None)
Anirudh's avatar
Anirudh committed
421

422
    model2 = model_fn(num_classes=50, memory_efficient=False)
Anirudh's avatar
Anirudh committed
423
424
425
426
427
428
429
    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)

430
431
432
    _check_input_backprop(model1, x)
    _check_input_backprop(model2, x)

Anirudh's avatar
Anirudh committed
433

434
435
436
@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
437
438
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
439
    model = models.resnet50(replace_stride_with_dilation=(dilate_layer_2, dilate_layer_3, dilate_layer_4))
Anirudh's avatar
Anirudh committed
440
441
442
443
444
445
446
447
448
    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():
449
    model = models.mobilenet_v2(inverted_residual_setting=[[1, 16, 1, 1], [6, 24, 2, 2]])
Anirudh's avatar
Anirudh committed
450
451
452
453
454
455
    model.eval()
    x = torch.rand(1, 3, 224, 224)
    out = model(x)
    assert out.shape[-1] == 1000


456
457
458
@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
459
460
461
    assert any(isinstance(x, nn.BatchNorm2d) for x in model.modules())

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

464
    model = model_fn(norm_layer=get_gn)
465
    assert not (any(isinstance(x, nn.BatchNorm2d) for x in model.modules()))
Anirudh's avatar
Anirudh committed
466
467
468
469
470
    assert any(isinstance(x, nn.GroupNorm) for x in model.modules())


def test_inception_v3_eval():
    kwargs = {}
471
472
473
    kwargs["transform_input"] = True
    kwargs["aux_logits"] = True
    kwargs["init_weights"] = False
Anirudh's avatar
Anirudh committed
474
475
476
477
478
479
480
    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))
481
    _check_input_backprop(model, x)
Anirudh's avatar
Anirudh committed
482
483
484


def test_fasterrcnn_double():
485
    model = models.detection.fasterrcnn_resnet50_fpn(num_classes=50, weights=None, weights_backbone=None)
Anirudh's avatar
Anirudh committed
486
487
488
489
490
491
492
493
494
495
496
    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]
497
    _check_input_backprop(model, model_input)
Anirudh's avatar
Anirudh committed
498
499
500
501


def test_googlenet_eval():
    kwargs = {}
502
503
504
    kwargs["transform_input"] = True
    kwargs["aux_logits"] = True
    kwargs["init_weights"] = False
Anirudh's avatar
Anirudh committed
505
506
507
508
509
510
511
512
    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))
513
    _check_input_backprop(model, x)
Anirudh's avatar
Anirudh committed
514
515
516
517
518
519
520
521
522
523


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

524
    model = models.detection.fasterrcnn_resnet50_fpn(num_classes=50, weights=None, weights_backbone=None)
Anirudh's avatar
Anirudh committed
525
526
527
    model.cuda()
    model.eval()
    input_shape = (3, 300, 300)
528
    x = torch.rand(input_shape, device="cuda")
Anirudh's avatar
Anirudh committed
529
530
531
532
533
534
535
    model_input = [x]
    out = model(model_input)
    assert model_input[0] is x

    checkOut(out)

    with torch.cuda.amp.autocast():
536
        out = model(model_input)
537

Anirudh's avatar
Anirudh committed
538
    checkOut(out)
539

540
541
    _check_input_backprop(model, model_input)

Anirudh's avatar
Anirudh committed
542
543
544
545
    # now switch to cpu and make sure it works
    model.cpu()
    x = x.cpu()
    out_cpu = model([x])
546

Anirudh's avatar
Anirudh committed
547
    checkOut(out_cpu)
548

549
550
    _check_input_backprop(model, [x])

551

Anirudh's avatar
Anirudh committed
552
def test_generalizedrcnn_transform_repr():
553

Anirudh's avatar
Anirudh committed
554
555
556
    min_size, max_size = 224, 299
    image_mean = [0.485, 0.456, 0.406]
    image_std = [0.229, 0.224, 0.225]
557

558
559
560
    t = models.detection.transform.GeneralizedRCNNTransform(
        min_size=min_size, max_size=max_size, image_mean=image_mean, image_std=image_std
    )
561

Anirudh's avatar
Anirudh committed
562
    # Check integrity of object __repr__ attribute
563
564
    expected_string = "GeneralizedRCNNTransform("
    _indent = "\n    "
565
566
    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
567
568
    expected_string += "mode='bilinear')\n)"
    assert t.__repr__() == expected_string
569
570


571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
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)


600
@pytest.mark.parametrize("model_fn", get_models_from_module(models))
601
@pytest.mark.parametrize("dev", cpu_and_gpu())
602
def test_classification_model(model_fn, dev):
Anirudh's avatar
Anirudh committed
603
604
    set_rng_seed(0)
    defaults = {
605
606
        "num_classes": 50,
        "input_shape": (1, 3, 224, 224),
Anirudh's avatar
Anirudh committed
607
    }
608
    model_name = model_fn.__name__
609
    if SKIP_BIG_MODEL and model_name in skipped_big_models:
610
        pytest.skip("Skipped to reduce memory usage. Set env var SKIP_BIG_MODEL=0 to enable test for this model")
Anirudh's avatar
Anirudh committed
611
    kwargs = {**defaults, **_model_params.get(model_name, {})}
612
    num_classes = kwargs.get("num_classes")
613
    input_shape = kwargs.pop("input_shape")
Anirudh's avatar
Anirudh committed
614

615
    model = model_fn(**kwargs)
Anirudh's avatar
Anirudh committed
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)
    _assert_expected(out.cpu(), model_name, prec=0.1)
621
    assert out.shape[-1] == num_classes
622
623
    _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
624

625
    if dev == "cuda":
Anirudh's avatar
Anirudh committed
626
627
628
629
630
631
        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
632

633
634
    _check_input_backprop(model, x)

635

636
@pytest.mark.parametrize("model_fn", get_models_from_module(models.segmentation))
637
@pytest.mark.parametrize("dev", cpu_and_gpu())
638
def test_segmentation_model(model_fn, dev):
Anirudh's avatar
Anirudh committed
639
640
    set_rng_seed(0)
    defaults = {
641
        "num_classes": 10,
642
        "weights_backbone": None,
643
        "input_shape": (1, 3, 32, 32),
Anirudh's avatar
Anirudh committed
644
    }
645
    model_name = model_fn.__name__
Anirudh's avatar
Anirudh committed
646
    kwargs = {**defaults, **_model_params.get(model_name, {})}
647
    input_shape = kwargs.pop("input_shape")
Anirudh's avatar
Anirudh committed
648

649
    model = model_fn(**kwargs)
Anirudh's avatar
Anirudh committed
650
651
652
    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)
653
    out = model(x)
Anirudh's avatar
Anirudh committed
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672

    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

673
    full_validation = check_out(out["out"])
Anirudh's avatar
Anirudh committed
674

675
676
    _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
677

678
    if dev == "cuda":
Anirudh's avatar
Anirudh committed
679
        with torch.cuda.amp.autocast():
680
            out = model(x)
Anirudh's avatar
Anirudh committed
681
682
            # See autocast_flaky_numerics comment at top of file.
            if model_name not in autocast_flaky_numerics:
683
                full_validation &= check_out(out["out"])
Anirudh's avatar
Anirudh committed
684
685

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

695
696
    _check_input_backprop(model, x)

697

698
@pytest.mark.parametrize("model_fn", get_models_from_module(models.detection))
699
@pytest.mark.parametrize("dev", cpu_and_gpu())
700
def test_detection_model(model_fn, dev):
Anirudh's avatar
Anirudh committed
701
702
    set_rng_seed(0)
    defaults = {
703
        "num_classes": 50,
704
        "weights_backbone": None,
705
        "input_shape": (3, 300, 300),
Anirudh's avatar
Anirudh committed
706
    }
707
    model_name = model_fn.__name__
Anirudh's avatar
Anirudh committed
708
    kwargs = {**defaults, **_model_params.get(model_name, {})}
709
    input_shape = kwargs.pop("input_shape")
Anirudh's avatar
Anirudh committed
710

711
    model = model_fn(**kwargs)
Anirudh's avatar
Anirudh committed
712
713
714
715
716
717
718
719
720
721
722
    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):
723
            tensor = tensor.cpu()
Anirudh's avatar
Anirudh committed
724
725
726
727
728
729
730
731
732
733
734
735
736
737
            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
738
            return tensor[ith_index - 1 :: ith_index]
Anirudh's avatar
Anirudh committed
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760

        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)
761
762
763
            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
764
765
766
767
768
769
770
771
772
773

            # 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)
774
    _check_jit_scriptable(model, ([x],), unwrapper=script_model_unwrapper.get(model_name, None), eager_out=out)
Anirudh's avatar
Anirudh committed
775

776
    if dev == "cuda":
Anirudh's avatar
Anirudh committed
777
778
779
780
781
782
783
        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:
784
        msg = (
785
            f"The output of {test_detection_model.__name__} could only be partially validated. "
786
787
            "This is likely due to unit-test flakiness, but you may "
            "want to do additional manual checks if you made "
788
            "significant changes to the codebase."
789
        )
Anirudh's avatar
Anirudh committed
790
791
        warnings.warn(msg, RuntimeWarning)
        pytest.skip(msg)
792

793
794
    _check_input_backprop(model, model_input)

795

796
797
@pytest.mark.parametrize("model_fn", get_models_from_module(models.detection))
def test_detection_model_validation(model_fn):
Anirudh's avatar
Anirudh committed
798
    set_rng_seed(0)
799
    model = model_fn(num_classes=50, weights=None, weights_backbone=None)
Anirudh's avatar
Anirudh committed
800
801
802
803
    input_shape = (3, 300, 300)
    x = [torch.rand(input_shape)]

    # validate that targets are present in training
804
    with pytest.raises(AssertionError):
Anirudh's avatar
Anirudh committed
805
806
807
        model(x)

    # validate type
808
    targets = [{"boxes": 0.0}]
809
    with pytest.raises(AssertionError):
Anirudh's avatar
Anirudh committed
810
811
812
813
        model(x, targets=targets)

    # validate boxes shape
    for boxes in (torch.rand((4,)), torch.rand((1, 5))):
814
        targets = [{"boxes": boxes}]
815
        with pytest.raises(AssertionError):
Anirudh's avatar
Anirudh committed
816
817
818
819
            model(x, targets=targets)

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

824

825
@pytest.mark.parametrize("model_fn", get_models_from_module(models.video))
826
@pytest.mark.parametrize("dev", cpu_and_gpu())
827
def test_video_model(model_fn, dev):
828
    set_rng_seed(0)
Anirudh's avatar
Anirudh committed
829
830
    # the default input shape is
    # bs * num_channels * clip_len * h *w
831
832
833
834
    defaults = {
        "input_shape": (1, 3, 4, 112, 112),
        "num_classes": 50,
    }
835
    model_name = model_fn.__name__
836
837
    if SKIP_BIG_MODEL and model_name in skipped_big_models:
        pytest.skip("Skipped to reduce memory usage. Set env var SKIP_BIG_MODEL=0 to enable test for this model")
838
839
840
    kwargs = {**defaults, **_model_params.get(model_name, {})}
    num_classes = kwargs.get("num_classes")
    input_shape = kwargs.pop("input_shape")
Anirudh's avatar
Anirudh committed
841
    # test both basicblock and Bottleneck
842
    model = model_fn(**kwargs)
Anirudh's avatar
Anirudh committed
843
844
845
846
    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)
847
848
    _assert_expected(out.cpu(), model_name, prec=0.1)
    assert out.shape[-1] == num_classes
849
850
    _check_jit_scriptable(model, (x,), unwrapper=script_model_unwrapper.get(model_name, None), eager_out=out)
    _check_fx_compatible(model, x, eager_out=out)
851
    assert out.shape[-1] == num_classes
Anirudh's avatar
Anirudh committed
852

853
    if dev == "cuda":
Anirudh's avatar
Anirudh committed
854
855
        with torch.cuda.amp.autocast():
            out = model(x)
856
857
858
859
            # 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] == num_classes
860

861
862
    _check_input_backprop(model, x)

863

864
865
866
867
868
869
870
@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",
)
871
872
@pytest.mark.parametrize("model_fn", get_models_from_module(models.quantization))
def test_quantized_classification_model(model_fn):
873
    set_rng_seed(0)
874
    defaults = {
875
        "num_classes": 5,
876
877
        "input_shape": (1, 3, 224, 224),
        "quantize": True,
878
    }
879
    model_name = model_fn.__name__
880
    kwargs = {**defaults, **_model_params.get(model_name, {})}
881
    input_shape = kwargs.pop("input_shape")
882
883

    # First check if quantize=True provides models that can run with input data
884
    model = model_fn(**kwargs)
885
    model.eval()
886
    x = torch.rand(input_shape)
887
888
889
890
891
    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
892
893
894
895
896
897
898
        _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
899

900
    kwargs["quantize"] = False
901
    for eval_mode in [True, False]:
902
        model = model_fn(**kwargs)
903
904
        if eval_mode:
            model.eval()
905
            model.qconfig = torch.ao.quantization.default_qconfig
906
907
        else:
            model.train()
908
            model.qconfig = torch.ao.quantization.default_qat_qconfig
909

910
        model.fuse_model(is_qat=not eval_mode)
911
        if eval_mode:
912
            torch.ao.quantization.prepare(model, inplace=True)
913
        else:
914
            torch.ao.quantization.prepare_qat(model, inplace=True)
915
916
            model.eval()

917
        torch.ao.quantization.convert(model, inplace=True)
918
919


920
@pytest.mark.parametrize("model_fn", get_models_from_module(models.detection))
921
def test_detection_model_trainable_backbone_layers(model_fn, disable_weight_loading):
922
    model_name = model_fn.__name__
923
924
925
    max_trainable = _model_tests_values[model_name]["max_trainable"]
    n_trainable_params = []
    for trainable_layers in range(0, max_trainable + 1):
926
        model = model_fn(weights=None, weights_backbone="DEFAULT", trainable_backbone_layers=trainable_layers)
927
928
929
930
931

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


932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
@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)


960
if __name__ == "__main__":
961
    pytest.main([__file__])