Function test_classification_model

test/test_models.py:678–716

(model_fn, dev)

@pytest.mark.parametrize("model_fn", list_model_fns(models))
@pytest.mark.parametrize("dev", cpu_and_cuda())
def test_classification_model(model_fn, dev):
    set_rng_seed(0)
    defaults = {
        "num_classes": 50,
        "input_shape": (1, 3, 224, 224),
    }
    model_name = model_fn.__name__
    if SKIP_BIG_MODEL and is_skippable(model_name, dev):
        pytest.skip("Skipped to reduce memory usage. Set env var SKIP_BIG_MODEL=0 to enable test for this model")
    kwargs = {**defaults, **_model_params.get(model_name, {})}
    num_classes = kwargs.get("num_classes")
    input_shape = kwargs.pop("input_shape")
    real_image = kwargs.pop("real_image", False)

    model = model_fn(**kwargs)
    model.eval().to(device=dev)
    x = _get_image(input_shape=input_shape, real_image=real_image, device=dev)
    out = model(x)
    # FIXME: this if/else is nasty and only here to please our CI prior to the
    # release. We rethink these tests altogether.
    if model_name == "resnet101":
        prec = 0.2
    else:
        # FIXME: this is probably still way too high.
        prec = 0.1
    _assert_expected(out.cpu(), model_name, prec=prec)
    assert out.shape[-1] == num_classes
    _check_jit_scriptable(model, (x,), unwrapper=script_model_unwrapper.get(model_name, None), eager_out=out)
    _check_fx_compatible(model, x, eager_out=out)

    if dev == "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

    _check_input_backprop(model, x)


@pytest.mark.parametrize("model_fn", list_model_fns(models.segmentation))