Port random affine, rotate, perspective and to_grayscale to pytest (#4000)

test/test_transforms_tensor.py

@@ -18,7 +18,7 @@ from common_utils import (
     _create_data_batch,
     _assert_equal_tensor_to_pil,
     _assert_approx_equal_tensor_to_pil,
-    cpu_and_gpu
+    cpu_and_gpu,
 )
 from _assert_utils import assert_equal
 
@@ -286,100 +286,6 @@ class Tester(unittest.TestCase):
         with get_tmp_dir() as tmp_dir:
             s_transform.save(os.path.join(tmp_dir, "t_resized_crop.pt"))
 
-    def test_random_affine(self):
-        tensor = torch.randint(0, 256, size=(3, 44, 56), dtype=torch.uint8, device=self.device)
-        batch_tensors = torch.randint(0, 256, size=(4, 3, 44, 56), dtype=torch.uint8, device=self.device)
-
-        def _test(**kwargs):
-            transform = T.RandomAffine(**kwargs)
-            s_transform = torch.jit.script(transform)
-
-            _test_transform_vs_scripted(transform, s_transform, tensor)
-            _test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors)
-
-            return s_transform
-
-        for interpolation in [NEAREST, BILINEAR]:
-            for shear in [15, 10.0, (5.0, 10.0), [-15, 15], [-10.0, 10.0, -11.0, 11.0]]:
-                _test(degrees=0.0, interpolation=interpolation, shear=shear)
-
-            for scale in [(0.7, 1.2), [0.7, 1.2]]:
-                _test(degrees=0.0, interpolation=interpolation, scale=scale)
-
-            for translate in [(0.1, 0.2), [0.2, 0.1]]:
-                _test(degrees=0.0, interpolation=interpolation, translate=translate)
-
-            for degrees in [45, 35.0, (-45, 45), [-90.0, 90.0]]:
-                _test(degrees=degrees, interpolation=interpolation)
-
-            for fill in [85, (10, -10, 10), 0.7, [0.0, 0.0, 0.0], [1, ], 1]:
-                _test(degrees=0.0, interpolation=interpolation, fill=fill)
-
-        s_transform = _test(degrees=0.0)
-        with get_tmp_dir() as tmp_dir:
-            s_transform.save(os.path.join(tmp_dir, "t_random_affine.pt"))
-
-    def test_random_rotate(self):
-        tensor = torch.randint(0, 256, size=(3, 44, 56), dtype=torch.uint8, device=self.device)
-        batch_tensors = torch.randint(0, 256, size=(4, 3, 44, 56), dtype=torch.uint8, device=self.device)
-
-        for center in [(0, 0), [10, 10], None, (56, 44)]:
-            for expand in [True, False]:
-                for degrees in [45, 35.0, (-45, 45), [-90.0, 90.0]]:
-                    for interpolation in [NEAREST, BILINEAR]:
-                        for fill in [85, (10, -10, 10), 0.7, [0.0, 0.0, 0.0], [1, ], 1]:
-                            transform = T.RandomRotation(
-                                degrees=degrees, interpolation=interpolation, expand=expand, center=center, fill=fill
-                            )
-                            s_transform = torch.jit.script(transform)
-
-                            _test_transform_vs_scripted(transform, s_transform, tensor)
-                            _test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors)
-
-        with get_tmp_dir() as tmp_dir:
-            s_transform.save(os.path.join(tmp_dir, "t_random_rotate.pt"))
-
-    def test_random_perspective(self):
-        tensor = torch.randint(0, 256, size=(3, 44, 56), dtype=torch.uint8, device=self.device)
-        batch_tensors = torch.randint(0, 256, size=(4, 3, 44, 56), dtype=torch.uint8, device=self.device)
-
-        for distortion_scale in np.linspace(0.1, 1.0, num=20):
-            for interpolation in [NEAREST, BILINEAR]:
-                for fill in [85, (10, -10, 10), 0.7, [0.0, 0.0, 0.0], [1, ], 1]:
-                    transform = T.RandomPerspective(
-                        distortion_scale=distortion_scale,
-                        interpolation=interpolation,
-                        fill=fill
-                    )
-                    s_transform = torch.jit.script(transform)
-
-                    _test_transform_vs_scripted(transform, s_transform, tensor)
-                    _test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors)
-
-        with get_tmp_dir() as tmp_dir:
-            s_transform.save(os.path.join(tmp_dir, "t_perspective.pt"))
-
-    def test_to_grayscale(self):
-
-        meth_kwargs = {"num_output_channels": 1}
-        tol = 1.0 + 1e-10
-        _test_class_op(
-            T.Grayscale, meth_kwargs=meth_kwargs, test_exact_match=False, device=self.device,
-            tol=tol, agg_method="max"
-        )
-
-        meth_kwargs = {"num_output_channels": 3}
-        _test_class_op(
-            T.Grayscale, meth_kwargs=meth_kwargs, test_exact_match=False, device=self.device,
-            tol=tol, agg_method="max"
-        )
-
-        meth_kwargs = {}
-        _test_class_op(
-            T.RandomGrayscale, meth_kwargs=meth_kwargs, test_exact_match=False, device=self.device,
-            tol=tol, agg_method="max"
-        )
-
     def test_normalize(self):
         fn = T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
         tensor, _ = _create_data(26, 34, device=self.device)
@@ -736,5 +642,125 @@ class CUDATester(Tester):
         self.device = "cuda"
 
 
+def _test_random_affine_helper(device, **kwargs):
+    tensor = torch.randint(0, 256, size=(3, 44, 56), dtype=torch.uint8, device=device)
+    batch_tensors = torch.randint(0, 256, size=(4, 3, 44, 56), dtype=torch.uint8, device=device)
+    transform = T.RandomAffine(**kwargs)
+    s_transform = torch.jit.script(transform)
+
+    _test_transform_vs_scripted(transform, s_transform, tensor)
+    _test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors)
+
+
+@pytest.mark.parametrize('device', cpu_and_gpu())
+def test_random_affine(device):
+    transform = T.RandomAffine(degrees=45.0)
+    s_transform = torch.jit.script(transform)
+    with get_tmp_dir() as tmp_dir:
+        s_transform.save(os.path.join(tmp_dir, "t_random_affine.pt"))
+
+
+@pytest.mark.parametrize('device', cpu_and_gpu())
+@pytest.mark.parametrize('interpolation', [NEAREST, BILINEAR])
+@pytest.mark.parametrize('shear', [15, 10.0, (5.0, 10.0), [-15, 15], [-10.0, 10.0, -11.0, 11.0]])
+def test_random_affine_shear(device, interpolation, shear):
+    _test_random_affine_helper(device, degrees=0.0, interpolation=interpolation, shear=shear)
+
+
+@pytest.mark.parametrize('device', cpu_and_gpu())
+@pytest.mark.parametrize('interpolation', [NEAREST, BILINEAR])
+@pytest.mark.parametrize('scale', [(0.7, 1.2), [0.7, 1.2]])
+def test_random_affine_scale(device, interpolation, scale):
+    _test_random_affine_helper(device, degrees=0.0, interpolation=interpolation, scale=scale)
+
+
+@pytest.mark.parametrize('device', cpu_and_gpu())
+@pytest.mark.parametrize('interpolation', [NEAREST, BILINEAR])
+@pytest.mark.parametrize('translate', [(0.1, 0.2), [0.2, 0.1]])
+def test_random_affine_translate(device, interpolation, translate):
+    _test_random_affine_helper(device, degrees=0.0, interpolation=interpolation, translate=translate)
+
+
+@pytest.mark.parametrize('device', cpu_and_gpu())
+@pytest.mark.parametrize('interpolation', [NEAREST, BILINEAR])
+@pytest.mark.parametrize('degrees', [45, 35.0, (-45, 45), [-90.0, 90.0]])
+def test_random_affine_degrees(device, interpolation, degrees):
+    _test_random_affine_helper(device, degrees=degrees, interpolation=interpolation)
+
+
+@pytest.mark.parametrize('device', cpu_and_gpu())
+@pytest.mark.parametrize('interpolation', [NEAREST, BILINEAR])
+@pytest.mark.parametrize('fill', [85, (10, -10, 10), 0.7, [0.0, 0.0, 0.0], [1, ], 1])
+def test_random_affine_fill(device, interpolation, fill):
+    _test_random_affine_helper(device, degrees=0.0, interpolation=interpolation, fill=fill)
+
+
+@pytest.mark.parametrize('device', cpu_and_gpu())
+@pytest.mark.parametrize('center', [(0, 0), [10, 10], None, (56, 44)])
+@pytest.mark.parametrize('expand', [True, False])
+@pytest.mark.parametrize('degrees', [45, 35.0, (-45, 45), [-90.0, 90.0]])
+@pytest.mark.parametrize('interpolation', [NEAREST, BILINEAR])
+@pytest.mark.parametrize('fill', [85, (10, -10, 10), 0.7, [0.0, 0.0, 0.0], [1, ], 1])
+def test_random_rotate(device, center, expand, degrees, interpolation, fill):
+    tensor = torch.randint(0, 256, size=(3, 44, 56), dtype=torch.uint8, device=device)
+    batch_tensors = torch.randint(0, 256, size=(4, 3, 44, 56), dtype=torch.uint8, device=device)
+
+    transform = T.RandomRotation(
+        degrees=degrees, interpolation=interpolation, expand=expand, center=center, fill=fill
+    )
+    s_transform = torch.jit.script(transform)
+
+    _test_transform_vs_scripted(transform, s_transform, tensor)
+    _test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors)
+
+
+def test_random_rotate_save():
+    transform = T.RandomRotation(degrees=45.0)
+    s_transform = torch.jit.script(transform)
+    with get_tmp_dir() as tmp_dir:
+        s_transform.save(os.path.join(tmp_dir, "t_random_rotate.pt"))
+
+
+@pytest.mark.parametrize('device', cpu_and_gpu())
+@pytest.mark.parametrize('distortion_scale', np.linspace(0.1, 1.0, num=20))
+@pytest.mark.parametrize('interpolation', [NEAREST, BILINEAR])
+@pytest.mark.parametrize('fill', [85, (10, -10, 10), 0.7, [0.0, 0.0, 0.0], [1, ], 1])
+def test_random_perspective(device, distortion_scale, interpolation, fill):
+    tensor = torch.randint(0, 256, size=(3, 44, 56), dtype=torch.uint8, device=device)
+    batch_tensors = torch.randint(0, 256, size=(4, 3, 44, 56), dtype=torch.uint8, device=device)
+
+    transform = T.RandomPerspective(
+        distortion_scale=distortion_scale,
+        interpolation=interpolation,
+        fill=fill
+    )
+    s_transform = torch.jit.script(transform)
+
+    _test_transform_vs_scripted(transform, s_transform, tensor)
+    _test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors)
+
+
+def test_random_perspective_save():
+    transform = T.RandomPerspective()
+    s_transform = torch.jit.script(transform)
+    with get_tmp_dir() as tmp_dir:
+        s_transform.save(os.path.join(tmp_dir, "t_perspective.pt"))
+
+
+@pytest.mark.parametrize('device', cpu_and_gpu())
+@pytest.mark.parametrize('Klass, meth_kwargs', [
+    (T.Grayscale, {"num_output_channels": 1}),
+    (T.Grayscale, {"num_output_channels": 3}),
+    (T.RandomGrayscale, {})
+])
+def test_to_grayscale(device, Klass, meth_kwargs):
+
+    tol = 1.0 + 1e-10
+    _test_class_op(
+        Klass, meth_kwargs=meth_kwargs, test_exact_match=False, device=device,
+        tol=tol, agg_method="max"
+    )
+
+
 if __name__ == '__main__':
     unittest.main()