simplify normalize testing into single test parameterize on input creation

Author: justincdavis

test/test_transforms_v2.py

@@ -5597,24 +5597,36 @@ def _reference_normalize_image(self, image, *, mean, std):
 
     @pytest.mark.parametrize(("mean", "std"), MEANS_STDS)
     @pytest.mark.parametrize("dtype", [torch.float16, torch.float32, torch.float64])
+    @pytest.mark.parametrize(
+        "make_input",
+        [
+            make_image,
+            pytest.param(
+                make_image_cvcuda, marks=pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="test requires CVCUDA")
+            ),
+        ],
+    )
     @pytest.mark.parametrize("fn", [F.normalize, transform_cls_to_functional(transforms.Normalize)])
-    def test_correctness_image(self, mean, std, dtype, fn):
-        image = make_image(dtype=dtype)
+    def test_correctness_image(self, mean, std, dtype, make_input, fn):
+        if make_input == make_image_cvcuda and dtype != torch.float32:
+            pytest.skip("CVCUDA only supports float32 for normalize")
+
+        image = make_input(dtype=dtype)
 
         actual = fn(image, mean=mean, std=std)
-        expected = self._reference_normalize_image(image, mean=mean, std=std)
 
-        assert_equal(actual, expected)
+        if make_input == make_image_cvcuda:
+            image = F.cvcuda_to_tensor(image).to(device="cpu")
+            image = image.squeeze(0)
+            actual = F.cvcuda_to_tensor(actual).to(device="cpu")
+            actual = actual.squeeze(0)
 
-    @pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="test requires CVCUDA")
-    @pytest.mark.parametrize(("mean", "std"), MEANS_STDS)
-    @pytest.mark.parametrize("fn", [F.normalize, transform_cls_to_functional(transforms.Normalize)])
-    def test_correctness_cvcuda(self, mean, std, fn):
-        image = make_image(batch_dims=(1,), dtype=torch.float32, device="cuda")
-        cvc_image = F.to_cvcuda_tensor(image)
-        actual = F._misc._normalize_cvcuda(cvc_image, mean=mean, std=std)
-        expected = fn(image, mean=mean, std=std)
-        torch.testing.assert_close(F.cvcuda_to_tensor(actual), expected, rtol=1e-7, atol=1e-7)
+        expected = self._reference_normalize_image(image, mean=mean, std=std)
+
+        if make_input == make_image_cvcuda:
+            torch.testing.assert_close(actual, expected, rtol=0, atol=1e-6)
+        else:
+            assert_equal(actual, expected)
 
 
 class TestClampBoundingBoxes:

torchvision/transforms/v2/_transform.py

@@ -8,7 +8,7 @@
 from torch import nn
 from torch.utils._pytree import tree_flatten, tree_unflatten
 from torchvision import tv_tensors
-from torchvision.transforms.v2._utils import check_type, has_any, is_pure_tensor
+from torchvision.transforms.v2._utils import check_type, has_any, is_cvcuda_tensor, is_pure_tensor
 from torchvision.utils import _log_api_usage_once
 
 from .functional._utils import _get_kernel
@@ -23,7 +23,7 @@ class Transform(nn.Module):
 
     # Class attribute defining transformed types. Other types are passed-through without any transformation
     # We support both Types and callables that are able to do further checks on the type of the input.
-    _transformed_types: tuple[type | Callable[[Any], bool], ...] = (torch.Tensor, PIL.Image.Image)
+    _transformed_types: tuple[type | Callable[[Any], bool], ...] = (torch.Tensor, PIL.Image.Image, is_cvcuda_tensor)
 
     def __init__(self) -> None:
         super().__init__()

torchvision/transforms/v2/_utils.py

@@ -15,7 +15,7 @@
 from torchvision._utils import sequence_to_str
 
 from torchvision.transforms.transforms import _check_sequence_input, _setup_angle, _setup_size  # noqa: F401
-from torchvision.transforms.v2.functional import get_dimensions, get_size, is_pure_tensor
+from torchvision.transforms.v2.functional import get_dimensions, get_size, is_cvcuda_tensor, is_pure_tensor
 from torchvision.transforms.v2.functional._utils import _FillType, _FillTypeJIT
 
 
@@ -182,7 +182,7 @@ def query_chw(flat_inputs: list[Any]) -> tuple[int, int, int]:
     chws = {
         tuple(get_dimensions(inpt))
         for inpt in flat_inputs
-        if check_type(inpt, (is_pure_tensor, tv_tensors.Image, PIL.Image.Image, tv_tensors.Video))
+        if check_type(inpt, (is_pure_tensor, tv_tensors.Image, PIL.Image.Image, tv_tensors.Video, is_cvcuda_tensor))
     }
     if not chws:
         raise TypeError("No image or video was found in the sample")
@@ -207,6 +207,7 @@ def query_size(flat_inputs: list[Any]) -> tuple[int, int]:
                 tv_tensors.Mask,
                 tv_tensors.BoundingBoxes,
                 tv_tensors.KeyPoints,
+                is_cvcuda_tensor,
             ),
         )
     }

torchvision/transforms/v2/functional/__init__.py

@@ -1,6 +1,6 @@
 from torchvision.transforms import InterpolationMode  # usort: skip
 
-from ._utils import is_pure_tensor, register_kernel  # usort: skip
+from ._utils import is_pure_tensor, register_kernel, is_cvcuda_tensor  # usort: skip
 
 from ._meta import (
     clamp_bounding_boxes,

torchvision/transforms/v2/functional/_utils.py

@@ -169,3 +169,10 @@ def _is_cvcuda_available():
         return True
     except ImportError:
         return False
+
+
+def is_cvcuda_tensor(inpt: Any) -> bool:
+    if _is_cvcuda_available():
+        cvcuda = _import_cvcuda()
+        return isinstance(inpt, cvcuda.Tensor)
+    return False