test_models.py 27.2 KB
Newer Older
1
import functools
2
import io
3
4
5
6
import operator
import os
import traceback
import warnings
7
from collections import OrderedDict
8
9

import pytest
10
import torch
11
import torch.fx
12
import torch.nn as nn
13
14
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
15
from torchvision import models
16

eellison's avatar
eellison committed
17

18
ACCEPT = os.getenv("EXPECTTEST_ACCEPT", "0") == "1"
19
20


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


26
27
28
29
30
31
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
32
    expected_file = expected_file = os.path.join(expected_file_base, "ModelTester.test_" + name)
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
    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


def _assert_expected(output, name, prec):
    """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)}
57
        print(f"Accepting updated output for {filename}:\n\n{output}")
58
59
60
61
        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:
62
            raise RuntimeError(f"The output for {filename}, is larger than 50kb")
63
64
65
66
67
68
69
70
71
72
73
    else:
        expected = torch.load(expected_file)
        rtol = atol = prec
        torch.testing.assert_close(output, expected, rtol=rtol, atol=atol, check_dtype=False)


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

    def assert_export_import_module(m, args):
        """Check that the results of a model are the same after saving and loading"""
74

75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
        def get_export_import_copy(m):
            """Save and load a TorchScript model"""
            buffer = io.BytesIO()
            torch.jit.save(m, buffer)
            buffer.seek(0)
            imported = torch.jit.load(buffer)
            return imported

        m_import = get_export_import_copy(m)
        with freeze_rng_state():
            results = m(*args)
        with freeze_rng_state():
            results_from_imported = m_import(*args)
        tol = 3e-4
        try:
            torch.testing.assert_close(results, results_from_imported, atol=tol, rtol=tol)
91
        except ValueError:
92
93
94
95
96
97
            # custom check for the models that return named tuples:
            # we compare field by field while ignoring None as assert_close can't handle None
            for a, b in zip(results, results_from_imported):
                if a is not None:
                    torch.testing.assert_close(a, b, atol=tol, rtol=tol)

98
    TEST_WITH_SLOW = os.getenv("PYTORCH_TEST_WITH_SLOW", "0") == "1"
99
100
    if not TEST_WITH_SLOW or skip:
        # TorchScript is not enabled, skip these tests
101
        msg = (
102
            f"The check_jit_scriptable test for {nn_module.__class__.__name__} was skipped. "
103
104
105
106
            "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 "
107
            "manually skipped."
108
        )
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
        warnings.warn(msg, RuntimeWarning)
        return None

    sm = torch.jit.script(nn_module)

    with freeze_rng_state():
        eager_out = nn_module(*args)

    with freeze_rng_state():
        script_out = sm(*args)
        if unwrapper:
            script_out = unwrapper(script_out)

    torch.testing.assert_close(eager_out, script_out, atol=1e-4, rtol=1e-4)
    assert_export_import_module(sm, args)


126
127
128
129
130
131
132
def _check_fx_compatible(model, inputs):
    model_fx = torch.fx.symbolic_trace(model)
    out = model(inputs)
    out_fx = model_fx(inputs)
    torch.testing.assert_close(out, out_fx)


133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
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)


162
163
164
# 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
165
script_model_unwrapper = {
166
167
    "googlenet": lambda x: x.logits,
    "inception_v3": lambda x: x.logits,
168
    "fasterrcnn_resnet50_fpn": lambda x: x[1],
169
    "fasterrcnn_mobilenet_v3_large_fpn": lambda x: x[1],
170
    "fasterrcnn_mobilenet_v3_large_320_fpn": lambda x: x[1],
171
172
173
    "maskrcnn_resnet50_fpn": lambda x: x[1],
    "keypointrcnn_resnet50_fpn": lambda x: x[1],
    "retinanet_resnet50_fpn": lambda x: x[1],
174
    "ssd300_vgg16": lambda x: x[1],
175
    "ssdlite320_mobilenet_v3_large": lambda x: x[1],
176
}
177
178


179
180
181
182
183
184
185
186
187
188
189
190
191
192
# 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",
193
194
    "deeplabv3_resnet50",
    "deeplabv3_resnet101",
195
    "deeplabv3_mobilenet_v3_large",
196
197
    "fcn_resnet50",
    "fcn_resnet101",
198
    "lraspp_mobilenet_v3_large",
199
    "maskrcnn_resnet50_fpn",
200
201
)

202
203
204
# 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.
205
quantized_flaky_models = ("inception_v3", "resnet50")
206

207

208
209
210
# The following contains configuration parameters for all models which are used by
# the _test_*_model methods.
_model_params = {
211
212
213
214
215
216
217
    "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),
218
    },
219
220
221
222
223
224
    "keypointrcnn_resnet50_fpn": {
        "num_classes": 2,
        "min_size": 224,
        "max_size": 224,
        "box_score_thresh": 0.15,
        "input_shape": (3, 224, 224),
225
    },
226
227
228
229
230
    "fasterrcnn_resnet50_fpn": {
        "num_classes": 20,
        "min_size": 224,
        "max_size": 224,
        "input_shape": (3, 224, 224),
231
    },
232
233
234
235
236
    "maskrcnn_resnet50_fpn": {
        "num_classes": 10,
        "min_size": 224,
        "max_size": 224,
        "input_shape": (3, 224, 224),
237
    },
238
239
    "fasterrcnn_mobilenet_v3_large_fpn": {
        "box_score_thresh": 0.02076,
240
    },
241
242
243
244
    "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,
245
246
247
248
    },
}


249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
# The following contains configuration and expected values to be used tests that are model specific
_model_tests_values = {
    "retinanet_resnet50_fpn": {
        "max_trainable": 5,
        "n_trn_params_per_layer": [36, 46, 65, 78, 88, 89],
    },
    "keypointrcnn_resnet50_fpn": {
        "max_trainable": 5,
        "n_trn_params_per_layer": [48, 58, 77, 90, 100, 101],
    },
    "fasterrcnn_resnet50_fpn": {
        "max_trainable": 5,
        "n_trn_params_per_layer": [30, 40, 59, 72, 82, 83],
    },
    "maskrcnn_resnet50_fpn": {
        "max_trainable": 5,
        "n_trn_params_per_layer": [42, 52, 71, 84, 94, 95],
    },
    "fasterrcnn_mobilenet_v3_large_fpn": {
        "max_trainable": 6,
        "n_trn_params_per_layer": [22, 23, 44, 70, 91, 97, 100],
    },
    "fasterrcnn_mobilenet_v3_large_320_fpn": {
        "max_trainable": 6,
        "n_trn_params_per_layer": [22, 23, 44, 70, 91, 97, 100],
    },
    "ssd300_vgg16": {
        "max_trainable": 5,
        "n_trn_params_per_layer": [45, 51, 57, 63, 67, 71],
    },
    "ssdlite320_mobilenet_v3_large": {
        "max_trainable": 6,
        "n_trn_params_per_layer": [96, 99, 138, 200, 239, 257, 266],
    },
}


Anirudh's avatar
Anirudh committed
286
287
288
289
290
291
292
293
294
295
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


296
297
@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
298
299
300
    input_shape = (1, 3, 300, 300)
    x = torch.rand(input_shape)

301
    model1 = model_fn(num_classes=50, memory_efficient=True)
Anirudh's avatar
Anirudh committed
302
    params = model1.state_dict()
303
    num_params = sum(x.numel() for x in model1.parameters())
Anirudh's avatar
Anirudh committed
304
305
306
    model1.eval()
    out1 = model1(x)
    out1.sum().backward()
307
    num_grad = sum(x.grad.numel() for x in model1.parameters() if x.grad is not None)
Anirudh's avatar
Anirudh committed
308

309
    model2 = model_fn(num_classes=50, memory_efficient=False)
Anirudh's avatar
Anirudh committed
310
311
312
313
314
315
316
    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)

317
318
319
    _check_input_backprop(model1, x)
    _check_input_backprop(model2, x)

Anirudh's avatar
Anirudh committed
320

321
322
323
@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
324
325
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
326
    model = models.resnet50(replace_stride_with_dilation=(dilate_layer_2, dilate_layer_3, dilate_layer_4))
Anirudh's avatar
Anirudh committed
327
328
329
330
331
332
333
334
335
    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():
336
    model = models.mobilenet_v2(inverted_residual_setting=[[1, 16, 1, 1], [6, 24, 2, 2]])
Anirudh's avatar
Anirudh committed
337
338
339
340
341
342
    model.eval()
    x = torch.rand(1, 3, 224, 224)
    out = model(x)
    assert out.shape[-1] == 1000


343
344
345
@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
346
347
348
349
350
    assert any(isinstance(x, nn.BatchNorm2d) for x in model.modules())

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

351
    model = model_fn(norm_layer=get_gn)
352
    assert not (any(isinstance(x, nn.BatchNorm2d) for x in model.modules()))
Anirudh's avatar
Anirudh committed
353
354
355
356
357
358
    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 = {}
359
360
361
    kwargs["transform_input"] = True
    kwargs["aux_logits"] = True
    kwargs["init_weights"] = False
Anirudh's avatar
Anirudh committed
362
363
364
365
366
367
368
    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))
369
    _check_input_backprop(model, x)
Anirudh's avatar
Anirudh committed
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384


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]
385
    _check_input_backprop(model, model_input)
Anirudh's avatar
Anirudh committed
386
387
388
389
390


def test_googlenet_eval():
    # replacement for models.googlenet(pretrained=True) that does not download weights
    kwargs = {}
391
392
393
    kwargs["transform_input"] = True
    kwargs["aux_logits"] = True
    kwargs["init_weights"] = False
Anirudh's avatar
Anirudh committed
394
395
396
397
398
399
400
401
    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))
402
    _check_input_backprop(model, x)
Anirudh's avatar
Anirudh committed
403
404
405
406
407
408
409
410
411
412
413
414
415
416


@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)
417
    x = torch.rand(input_shape, device="cuda")
Anirudh's avatar
Anirudh committed
418
419
420
421
422
423
424
    model_input = [x]
    out = model(model_input)
    assert model_input[0] is x

    checkOut(out)

    with torch.cuda.amp.autocast():
425
        out = model(model_input)
426

Anirudh's avatar
Anirudh committed
427
    checkOut(out)
428

429
430
    _check_input_backprop(model, model_input)

Anirudh's avatar
Anirudh committed
431
432
433
434
    # now switch to cpu and make sure it works
    model.cpu()
    x = x.cpu()
    out_cpu = model([x])
435

Anirudh's avatar
Anirudh committed
436
    checkOut(out_cpu)
437

438
439
    _check_input_backprop(model, [x])

440

Anirudh's avatar
Anirudh committed
441
def test_generalizedrcnn_transform_repr():
442

Anirudh's avatar
Anirudh committed
443
444
445
    min_size, max_size = 224, 299
    image_mean = [0.485, 0.456, 0.406]
    image_std = [0.229, 0.224, 0.225]
446

447
448
449
    t = models.detection.transform.GeneralizedRCNNTransform(
        min_size=min_size, max_size=max_size, image_mean=image_mean, image_std=image_std
    )
450

Anirudh's avatar
Anirudh committed
451
    # Check integrity of object __repr__ attribute
452
453
    expected_string = "GeneralizedRCNNTransform("
    _indent = "\n    "
454
455
    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
456
457
    expected_string += "mode='bilinear')\n)"
    assert t.__repr__() == expected_string
458
459


460
@pytest.mark.parametrize("model_fn", get_models_from_module(models))
461
@pytest.mark.parametrize("dev", cpu_and_gpu())
462
def test_classification_model(model_fn, dev):
Anirudh's avatar
Anirudh committed
463
464
    set_rng_seed(0)
    defaults = {
465
466
        "num_classes": 50,
        "input_shape": (1, 3, 224, 224),
Anirudh's avatar
Anirudh committed
467
    }
468
    model_name = model_fn.__name__
Anirudh's avatar
Anirudh committed
469
    kwargs = {**defaults, **_model_params.get(model_name, {})}
470
    input_shape = kwargs.pop("input_shape")
Anirudh's avatar
Anirudh committed
471

472
    model = model_fn(**kwargs)
Anirudh's avatar
Anirudh committed
473
474
475
476
477
478
479
    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)
    assert out.shape[-1] == 50
    _check_jit_scriptable(model, (x,), unwrapper=script_model_unwrapper.get(model_name, None))
480
    _check_fx_compatible(model, x)
Anirudh's avatar
Anirudh committed
481
482
483
484
485
486
487
488

    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
489

490
491
    _check_input_backprop(model, x)

492

493
@pytest.mark.parametrize("model_fn", get_models_from_module(models.segmentation))
494
@pytest.mark.parametrize("dev", cpu_and_gpu())
495
def test_segmentation_model(model_fn, dev):
Anirudh's avatar
Anirudh committed
496
497
    set_rng_seed(0)
    defaults = {
498
499
500
        "num_classes": 10,
        "pretrained_backbone": False,
        "input_shape": (1, 3, 32, 32),
Anirudh's avatar
Anirudh committed
501
    }
502
    model_name = model_fn.__name__
Anirudh's avatar
Anirudh committed
503
    kwargs = {**defaults, **_model_params.get(model_name, {})}
504
    input_shape = kwargs.pop("input_shape")
Anirudh's avatar
Anirudh committed
505

506
    model = model_fn(**kwargs)
Anirudh's avatar
Anirudh committed
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
    model.eval().to(device=dev)
    # RNG always on CPU, to ensure x in cuda tests is bitwise identical to x in cpu tests
    x = torch.rand(input_shape).to(device=dev)
    out = model(x)["out"]

    def check_out(out):
        prec = 0.01
        try:
            # We first try to assert the entire output if possible. This is not
            # only the best way to assert results but also handles the cases
            # where we need to create a new expected result.
            _assert_expected(out.cpu(), model_name, prec=prec)
        except AssertionError:
            # Unfortunately some segmentation models are flaky with autocast
            # so instead of validating the probability scores, check that the class
            # predictions match.
            expected_file = _get_expected_file(model_name)
            expected = torch.load(expected_file)
            torch.testing.assert_close(out.argmax(dim=1), expected.argmax(dim=1), rtol=prec, atol=prec)
            return False  # Partial validation performed

        return True  # Full validation performed

    full_validation = check_out(out)

    _check_jit_scriptable(model, (x,), unwrapper=script_model_unwrapper.get(model_name, None))
533
    _check_fx_compatible(model, x)
Anirudh's avatar
Anirudh committed
534
535
536
537
538
539
540
541
542

    if dev == torch.device("cuda"):
        with torch.cuda.amp.autocast():
            out = model(x)["out"]
            # See autocast_flaky_numerics comment at top of file.
            if model_name not in autocast_flaky_numerics:
                full_validation &= check_out(out)

    if not full_validation:
543
        msg = (
544
            f"The output of {test_segmentation_model.__name__} could only be partially validated. "
545
546
            "This is likely due to unit-test flakiness, but you may "
            "want to do additional manual checks if you made "
547
            "significant changes to the codebase."
548
        )
Anirudh's avatar
Anirudh committed
549
550
        warnings.warn(msg, RuntimeWarning)
        pytest.skip(msg)
551

552
553
    _check_input_backprop(model, x)

554

555
@pytest.mark.parametrize("model_fn", get_models_from_module(models.detection))
556
@pytest.mark.parametrize("dev", cpu_and_gpu())
557
def test_detection_model(model_fn, dev):
Anirudh's avatar
Anirudh committed
558
559
    set_rng_seed(0)
    defaults = {
560
561
562
        "num_classes": 50,
        "pretrained_backbone": False,
        "input_shape": (3, 300, 300),
Anirudh's avatar
Anirudh committed
563
    }
564
    model_name = model_fn.__name__
Anirudh's avatar
Anirudh committed
565
    kwargs = {**defaults, **_model_params.get(model_name, {})}
566
    input_shape = kwargs.pop("input_shape")
Anirudh's avatar
Anirudh committed
567

568
    model = model_fn(**kwargs)
Anirudh's avatar
Anirudh committed
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
    model.eval().to(device=dev)
    # RNG always on CPU, to ensure x in cuda tests is bitwise identical to x in cpu tests
    x = torch.rand(input_shape).to(device=dev)
    model_input = [x]
    out = model(model_input)
    assert model_input[0] is x

    def check_out(out):
        assert len(out) == 1

        def compact(tensor):
            size = tensor.size()
            elements_per_sample = functools.reduce(operator.mul, size[1:], 1)
            if elements_per_sample > 30:
                return compute_mean_std(tensor)
            else:
                return subsample_tensor(tensor)

        def subsample_tensor(tensor):
            num_elems = tensor.size(0)
            num_samples = 20
            if num_elems <= num_samples:
                return tensor

            ith_index = num_elems // num_samples
594
            return tensor[ith_index - 1 :: ith_index]
Anirudh's avatar
Anirudh committed
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616

        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)
617
618
619
            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
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639

            # Note: Fmassa proposed turning off NMS by adapting the threshold
            # and then using the Hungarian algorithm as in DETR to find the
            # best match between output and expected boxes and eliminate some
            # of the flakiness. Worth exploring.
            return False  # Partial validation performed

        return True  # Full validation performed

    full_validation = check_out(out)
    _check_jit_scriptable(model, ([x],), unwrapper=script_model_unwrapper.get(model_name, None))

    if dev == torch.device("cuda"):
        with torch.cuda.amp.autocast():
            out = model(model_input)
            # See autocast_flaky_numerics comment at top of file.
            if model_name not in autocast_flaky_numerics:
                full_validation &= check_out(out)

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

649
650
    _check_input_backprop(model, model_input)

651

652
653
@pytest.mark.parametrize("model_fn", get_models_from_module(models.detection))
def test_detection_model_validation(model_fn):
Anirudh's avatar
Anirudh committed
654
    set_rng_seed(0)
655
    model = model_fn(num_classes=50, pretrained_backbone=False)
Anirudh's avatar
Anirudh committed
656
657
658
659
660
661
662
663
    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
664
    targets = [{"boxes": 0.0}]
Anirudh's avatar
Anirudh committed
665
666
667
668
669
    with pytest.raises(ValueError):
        model(x, targets=targets)

    # validate boxes shape
    for boxes in (torch.rand((4,)), torch.rand((1, 5))):
670
        targets = [{"boxes": boxes}]
Anirudh's avatar
Anirudh committed
671
672
673
674
675
        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]])
676
    targets = [{"boxes": boxes}]
Anirudh's avatar
Anirudh committed
677
678
    with pytest.raises(ValueError):
        model(x, targets=targets)
679

680

681
@pytest.mark.parametrize("model_fn", get_models_from_module(models.video))
682
@pytest.mark.parametrize("dev", cpu_and_gpu())
683
def test_video_model(model_fn, dev):
Anirudh's avatar
Anirudh committed
684
685
686
    # the default input shape is
    # bs * num_channels * clip_len * h *w
    input_shape = (1, 3, 4, 112, 112)
687
    model_name = model_fn.__name__
Anirudh's avatar
Anirudh committed
688
    # test both basicblock and Bottleneck
689
    model = model_fn(num_classes=50)
Anirudh's avatar
Anirudh committed
690
691
692
693
694
    model.eval().to(device=dev)
    # RNG always on CPU, to ensure x in cuda tests is bitwise identical to x in cpu tests
    x = torch.rand(input_shape).to(device=dev)
    out = model(x)
    _check_jit_scriptable(model, (x,), unwrapper=script_model_unwrapper.get(model_name, None))
695
    _check_fx_compatible(model, x)
Anirudh's avatar
Anirudh committed
696
697
698
699
700
701
    assert out.shape[-1] == 50

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

703
704
    _check_input_backprop(model, x)

705

706
707
708
709
710
711
712
@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",
)
713
714
@pytest.mark.parametrize("model_fn", get_models_from_module(models.quantization))
def test_quantized_classification_model(model_fn):
715
    set_rng_seed(0)
716
    defaults = {
717
        "num_classes": 5,
718
719
720
        "input_shape": (1, 3, 224, 224),
        "pretrained": False,
        "quantize": True,
721
    }
722
    model_name = model_fn.__name__
723
    kwargs = {**defaults, **_model_params.get(model_name, {})}
724
    input_shape = kwargs.pop("input_shape")
725
726

    # First check if quantize=True provides models that can run with input data
727
    model = model_fn(**kwargs)
728
    model.eval()
729
    x = torch.rand(input_shape)
730
731
732
733
734
735
736
    out = model(x)

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

738
    kwargs["quantize"] = False
739
    for eval_mode in [True, False]:
740
        model = model_fn(**kwargs)
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
        if eval_mode:
            model.eval()
            model.qconfig = torch.quantization.default_qconfig
        else:
            model.train()
            model.qconfig = torch.quantization.default_qat_qconfig

        model.fuse_model()
        if eval_mode:
            torch.quantization.prepare(model, inplace=True)
        else:
            torch.quantization.prepare_qat(model, inplace=True)
            model.eval()

        torch.quantization.convert(model, inplace=True)

    try:
        torch.jit.script(model)
    except Exception as e:
        tb = traceback.format_exc()
        raise AssertionError(f"model cannot be scripted. Traceback = {str(tb)}") from e


764
765
766
@pytest.mark.parametrize("model_fn", get_models_from_module(models.detection))
def test_detection_model_trainable_backbone_layers(model_fn):
    model_name = model_fn.__name__
767
768
769
    max_trainable = _model_tests_values[model_name]["max_trainable"]
    n_trainable_params = []
    for trainable_layers in range(0, max_trainable + 1):
770
        model = model_fn(pretrained=False, pretrained_backbone=True, trainable_backbone_layers=trainable_layers)
771
772
773
774
775

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


776
if __name__ == "__main__":
777
    pytest.main([__file__])