diff --git a/test/expect/ModelTester.test_mc3_18_expect.pkl b/test/expect/ModelTester.test_mc3_18_expect.pkl new file mode 100644 index 00000000000..938c5216050 Binary files /dev/null and b/test/expect/ModelTester.test_mc3_18_expect.pkl differ diff --git a/test/expect/ModelTester.test_r2plus1d_18_expect.pkl b/test/expect/ModelTester.test_r2plus1d_18_expect.pkl new file mode 100644 index 00000000000..b62a6d04af9 Binary files /dev/null and b/test/expect/ModelTester.test_r2plus1d_18_expect.pkl differ diff --git a/test/expect/ModelTester.test_r3d_18_expect.pkl b/test/expect/ModelTester.test_r3d_18_expect.pkl new file mode 100644 index 00000000000..11cf06bb1fc Binary files /dev/null and b/test/expect/ModelTester.test_r3d_18_expect.pkl differ diff --git a/test/test_models.py b/test/test_models.py index 7d29dde0678..0acef4dcef6 100644 --- a/test/test_models.py +++ b/test/test_models.py @@ -822,24 +822,36 @@ def test_detection_model_validation(model_fn): @pytest.mark.parametrize("model_fn", get_models_from_module(models.video)) @pytest.mark.parametrize("dev", cpu_and_gpu()) def test_video_model(model_fn, dev): + set_rng_seed(0) # the default input shape is # bs * num_channels * clip_len * h *w - input_shape = (1, 3, 4, 112, 112) + defaults = { + "input_shape": (1, 3, 4, 112, 112), + "num_classes": 50, + } model_name = model_fn.__name__ + kwargs = {**defaults, **_model_params.get(model_name, {})} + num_classes = kwargs.get("num_classes") + input_shape = kwargs.pop("input_shape") # test both basicblock and Bottleneck - model = model_fn(num_classes=50) + model = model_fn(**kwargs) 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] == 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) - assert out.shape[-1] == 50 + assert out.shape[-1] == num_classes if dev == "cuda": with torch.cuda.amp.autocast(): out = model(x) - assert out.shape[-1] == 50 + # 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 _check_input_backprop(model, x)