Use real weight and image for classification model test and relaxing precision requirement for general model tests

test/test_models.py

@@ -18,7 +18,7 @@
 from common_utils import cpu_and_gpu, freeze_rng_state, map_nested_tensor_object, needs_cuda, set_rng_seed
 from PIL import Image
 from torchvision import models, transforms
-from torchvision.models import get_model_builder, list_models
+from torchvision.models import get_model_builder, get_model_weights, get_weight, list_models
 
 
 ACCEPT = os.getenv("EXPECTTEST_ACCEPT", "0") == "1"
@@ -29,20 +29,8 @@ def list_model_fns(module):
     return [get_model_builder(name) for name in list_models(module)]
 
 
-def _get_image(input_shape, real_image, device):
-    """This routine loads a real or random image based on `real_image` argument.
-    Currently, the real image is utilized for the following list of models:
-    - `retinanet_resnet50_fpn`,
-    - `retinanet_resnet50_fpn_v2`,
-    - `keypointrcnn_resnet50_fpn`,
-    - `fasterrcnn_resnet50_fpn`,
-    - `fasterrcnn_resnet50_fpn_v2`,
-    - `fcos_resnet50_fpn`,
-    - `maskrcnn_resnet50_fpn`,
-    - `maskrcnn_resnet50_fpn_v2`,
-    in `test_classification_model` and `test_detection_model`.
-    To do so, a keyword argument `real_image` was added to the abovelisted models in `_model_params`
-    """
+def _get_image(input_shape, real_image, device, weights=None, dtype=None):
+    """This routine loads a real or random image based on `real_image` argument."""
     if real_image:
         # TODO: Maybe unify file discovery logic with test_image.py
         GRACE_HOPPER = os.path.join(
@@ -51,19 +39,26 @@ def _get_image(input_shape, real_image, device):
 
         img = Image.open(GRACE_HOPPER)
 
-        original_width, original_height = img.size
-
-        # make the image square
-        img = img.crop((0, 0, original_width, original_width))
-        img = img.resize(input_shape[1:3])
+        if weights is None:
+            original_width, original_height = img.size
+            # make the image square
+            img = img.crop((0, 0, original_width, original_width))
+            img = img.resize(input_shape[-2:])
 
-        convert_tensor = transforms.ToTensor()
-        image = convert_tensor(img)
+            convert_tensor = transforms.ToTensor()
+            image = convert_tensor(img)
+        else:
+            H, W = input_shape[-2:]
+            min_side = min(H, W)
+            preprocess = weights.transforms(resize_size=min_side, crop_size=min_side)
+            image = preprocess(img)
+        if len(input_shape) > len(image.size()):
+            image = image.unsqueeze(0)
         assert tuple(image.size()) == input_shape
-        return image.to(device=device)
+        return image.to(device=device, dtype=dtype)
 
     # RNG always on CPU, to ensure x in cuda tests is bitwise identical to x in cpu tests
-    return torch.rand(input_shape).to(device=device)
+    return torch.rand(input_shape).to(device=device, dtype=dtype)
 
 
 @pytest.fixture
@@ -195,7 +190,7 @@ def _check_fx_compatible(model, inputs, eager_out=None):
         eager_out = model(inputs)
     with torch.no_grad(), freeze_rng_state():
         fx_out = model_fx(inputs)
-    torch.testing.assert_close(eager_out, fx_out)
+    torch.testing.assert_close(eager_out, fx_out, atol=5e-5, rtol=5e-5)
 
 
 def _check_input_backprop(model, inputs):
@@ -278,11 +273,15 @@ def _check_input_backprop(model, inputs):
 # tests under test_quantized_classification_model will be skipped for the following models.
 quantized_flaky_models = ("inception_v3", "resnet50")
 
+# The tests for the following detection models are flaky due to precision of float32
+# we will do the test in float64 for these models
+detection_flaky_models = ("keypointrcnn_resnet50_fpn", "maskrcnn_resnet50_fpn_v2")
+
 
 # The following contains configuration parameters for all models which are used by
 # the _test_*_model methods.
 _model_params = {
-    "inception_v3": {"input_shape": (1, 3, 299, 299), "init_weights": True},
+    "inception_v3": {"input_shape": (1, 3, 299, 299)},
     "retinanet_resnet50_fpn": {
         "num_classes": 20,
         "score_thresh": 0.01,
@@ -354,6 +353,7 @@ def _check_input_backprop(model, inputs):
     "vit_h_14": {
         "image_size": 56,
         "input_shape": (1, 3, 56, 56),
+        "weight_name": None,
     },
     "mvit_v1_b": {
         "input_shape": (1, 3, 16, 224, 224),
@@ -364,7 +364,8 @@ def _check_input_backprop(model, inputs):
     "s3d": {
         "input_shape": (1, 3, 16, 224, 224),
     },
-    "googlenet": {"init_weights": True},
+    "regnet_y_128gf": {"weight_name": "IMAGENET1K_SWAG_LINEAR_V1"},
+    "vitc_b_16": {"weight_name": None},
 }
 # speeding up slow models:
 slow_models = [
@@ -390,7 +391,7 @@ def _check_input_backprop(model, inputs):
     "swin_v2_b",
 ]
 for m in slow_models:
-    _model_params[m] = {"input_shape": (1, 3, 64, 64)}
+    _model_params[m] = dict(_model_params.get(m, dict()), **{"input_shape": (1, 3, 64, 64)})
 
 
 # skip big models to reduce memory usage on CI test. We can exclude combinations of (platform-system, device).
@@ -648,6 +649,7 @@ def test_generalizedrcnn_transform_repr():
 
 
 def vitc_b_16(**kwargs: Any):
+    kwargs.pop("weights", None)
     return models.VisionTransformer(
         image_size=224,
         patch_size=16,
@@ -671,33 +673,46 @@ def test_vitc_models(model_fn, dev):
 def test_classification_model(model_fn, dev):
     set_rng_seed(0)
     defaults = {
-        "num_classes": 50,
+        "num_classes": 1000,
         "input_shape": (1, 3, 224, 224),
+        "num_classes_to_check": 50,
+        "real_image": True,
     }
     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")
+    num_classes_to_check = kwargs.pop("num_classes_to_check")
     input_shape = kwargs.pop("input_shape")
     real_image = kwargs.pop("real_image", False)
+    weight_name = kwargs.pop("weight_name", "IMAGENET1K_V1")
+    weight = None
+    if weight_name is not None:
+        weight_enum = get_model_weights(model_name)
+        weight = get_weight(f"{weight_enum.__name__}.{weight_name}")
 
-    model = model_fn(**kwargs)
+    model = model_fn(weights=weight, **kwargs)
     model.eval().to(device=dev)
-    x = _get_image(input_shape=input_shape, real_image=real_image, device=dev)
-    out = model(x)
-    _assert_expected(out.cpu(), model_name, prec=1e-3)
+    x = _get_image(input_shape=input_shape, real_image=real_image, device=dev, weights=weight)
+    with torch.no_grad(), freeze_rng_state():
+        out = model(x)
+        expect_out = out[:, :num_classes_to_check]
+    _assert_expected(expect_out.cpu(), model_name, prec=3e-2)
     assert out.shape[-1] == num_classes
+    assert expect_out.shape[-1] == num_classes_to_check
     _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():
+        with torch.cuda.amp.autocast(), torch.no_grad(), freeze_rng_state():
+            model.to(x.device)
             out = model(x)
+            expect_out = out[:, :num_classes_to_check]
             # 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
+                _assert_expected(expect_out.cpu(), model_name, prec=0.1)
+            assert expect_out.shape[-1] == num_classes_to_check
 
     _check_input_backprop(model, x)
 
@@ -777,13 +792,17 @@ def test_detection_model(model_fn, dev):
         "input_shape": (3, 300, 300),
     }
     model_name = model_fn.__name__
+    if model_name in detection_flaky_models:
+        dtype = torch.float64
+    else:
+        dtype = torch.get_default_dtype()
     kwargs = {**defaults, **_model_params.get(model_name, {})}
     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)
+    model.eval().to(device=dev, dtype=dtype)
+    x = _get_image(input_shape=input_shape, real_image=real_image, device=dev, dtype=dtype)
     model_input = [x]
     with torch.no_grad(), freeze_rng_state():
         out = model(model_input)
@@ -818,7 +837,7 @@ def compute_mean_std(tensor):
             return {"mean": mean, "std": std}
 
         output = map_nested_tensor_object(out, tensor_map_fn=compact)
-        prec = 0.01
+        prec = 3e-2
         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
@@ -917,7 +936,7 @@ def test_video_model(model_fn, 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=1e-5)
+    _assert_expected(out.cpu(), model_name, prec=3e-3)
     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)