Speed up Transformations tests (#3936)

test/test_transforms.py

@@ -1604,7 +1604,7 @@ class Tester(unittest.TestCase):
             # Accept 3 wrong pixels
             self.assertLess(n_diff_pixels, 3,
                             "a={}, t={}, s={}, sh={}\n".format(a, t, s, sh) +
-                            "n diff pixels={}\n".format(np.sum(np.array(result)[:, :, 0] != true_result[:, :, 0])))
+                            "n diff pixels={}\n".format(n_diff_pixels))
 
         # Test rotation
         a = 45
@@ -1623,9 +1623,9 @@ class Tester(unittest.TestCase):
         _test_transformation(a=0.0, t=(0.0, 0.0), s=1.0, sh=sh)
 
         # Test rotation, scale, translation, shear
-        for a in range(-90, 90, 25):
+        for a in range(-90, 90, 36):
             for t1 in range(-10, 10, 5):
-                for s in [0.75, 0.98, 1.0, 1.2, 1.4]:
+                for s in [0.77, 1.0, 1.27]:
                     for sh in range(-15, 15, 5):
                         _test_transformation(a=a, t=(t1, t1), s=s, sh=(sh, sh))
 

test/test_transforms_tensor.py

@@ -398,21 +398,32 @@ class Tester(TransformsTester):
         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 shear in [15, 10.0, (5.0, 10.0), [-15, 15], [-10.0, 10.0, -11.0, 11.0]]:
+        def _test(**kwargs):
+            transform = T.RandomAffine(**kwargs)
+            s_transform = torch.jit.script(transform)
+
+            self._test_transform_vs_scripted(transform, s_transform, tensor)
+            self._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]]:
-                for translate in [(0.1, 0.2), [0.2, 0.1]]:
-                    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.RandomAffine(
-                                    degrees=degrees, translate=translate,
-                                    scale=scale, shear=shear, interpolation=interpolation, fill=fill
-                                )
-                                s_transform = torch.jit.script(transform)
-
-                                self._test_transform_vs_scripted(transform, s_transform, tensor)
-                                self._test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors)
+                _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"))
 
@@ -663,7 +674,7 @@ class Tester(TransformsTester):
             for fill in [None, 85, (10, -10, 10), 0.7, [0.0, 0.0, 0.0], [1, ], 1]:
                 transform = T.AutoAugment(policy=policy, fill=fill)
                 s_transform = torch.jit.script(transform)
-                for _ in range(100):
+                for _ in range(25):
                     self._test_transform_vs_scripted(transform, s_transform, tensor)
                     self._test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors)