Adapted almost all functional tensor tests on CPU/CUDA

- fixed bug with transforms using generated grid
- remains *_crop, blocked by pytorch#2568
- TODO: test_adjustments

Author: vfdev-5

test/test_functional_tensor.py

@@ -15,18 +15,31 @@
 import torchvision.transforms.functional as F
 
 
+def run_on_cpu_and_cuda(fn):
+
+    devices = ["cpu", ]
+    if torch.cuda.is_available():
+        devices.append("cuda")
+
+    def wrapper(self, *args, **kwargs):
+        for device in devices:
+            fn(self, device, *args, **kwargs)
+
+    return wrapper
+
+
 class Tester(unittest.TestCase):
 
-    def _create_data(self, height=3, width=3, channels=3):
-        tensor = torch.randint(0, 255, (channels, height, width), dtype=torch.uint8)
-        pil_img = Image.fromarray(tensor.permute(1, 2, 0).contiguous().numpy())
+    def _create_data(self, height=3, width=3, channels=3, device="cpu"):
+        tensor = torch.randint(0, 255, (channels, height, width), dtype=torch.uint8, device=device)
+        pil_img = Image.fromarray(tensor.permute(1, 2, 0).contiguous().cpu().numpy())
         return tensor, pil_img
 
     def compareTensorToPIL(self, tensor, pil_image, msg=None):
         pil_tensor = torch.as_tensor(np.array(pil_image).transpose((2, 0, 1)))
         if msg is None:
             msg = "tensor:\n{} \ndid not equal PIL tensor:\n{}".format(tensor, pil_tensor)
-        self.assertTrue(tensor.equal(pil_tensor), msg)
+        self.assertTrue(tensor.cpu().equal(pil_tensor), msg)
 
     def approxEqualTensorToPIL(self, tensor, pil_image, tol=1e-5, msg=None):
         pil_tensor = torch.as_tensor(np.array(pil_image).transpose((2, 0, 1))).to(tensor)
@@ -36,9 +49,10 @@ def approxEqualTensorToPIL(self, tensor, pil_image, tol=1e-5, msg=None):
             msg="{}: mae={}, tol={}: \n{}\nvs\n{}".format(msg, mae, tol, tensor[0, :10, :10], pil_tensor[0, :10, :10])
         )
 
-    def test_vflip(self):
+    @run_on_cpu_and_cuda
+    def test_vflip(self, device):
         script_vflip = torch.jit.script(F_t.vflip)
-        img_tensor = torch.randn(3, 16, 16)
+        img_tensor = torch.randn(3, 16, 16).to(device=device)
         img_tensor_clone = img_tensor.clone()
         vflipped_img = F_t.vflip(img_tensor)
         vflipped_img_again = F_t.vflip(vflipped_img)
@@ -49,9 +63,10 @@ def test_vflip(self):
         vflipped_img_script = script_vflip(img_tensor)
         self.assertTrue(torch.equal(vflipped_img, vflipped_img_script))
 
-    def test_hflip(self):
+    @run_on_cpu_and_cuda
+    def test_hflip(self, device):
         script_hflip = torch.jit.script(F_t.hflip)
-        img_tensor = torch.randn(3, 16, 16)
+        img_tensor = torch.randn(3, 16, 16).to(device)
         img_tensor_clone = img_tensor.clone()
         hflipped_img = F_t.hflip(img_tensor)
         hflipped_img_again = F_t.hflip(hflipped_img)
@@ -62,9 +77,10 @@ def test_hflip(self):
         hflipped_img_script = script_hflip(img_tensor)
         self.assertTrue(torch.equal(hflipped_img, hflipped_img_script))
 
-    def test_crop(self):
+    @run_on_cpu_and_cuda
+    def test_crop(self, device):
         script_crop = torch.jit.script(F_t.crop)
-        img_tensor = torch.randint(0, 255, (3, 16, 16), dtype=torch.uint8)
+        img_tensor = torch.randint(0, 255, (3, 16, 16), dtype=torch.uint8, device=device)
         img_tensor_clone = img_tensor.clone()
         top = random.randint(0, 15)
         left = random.randint(0, 15)
@@ -73,7 +89,7 @@ def test_crop(self):
         img_cropped = F_t.crop(img_tensor, top, left, height, width)
         img_PIL = transforms.ToPILImage()(img_tensor)
         img_PIL_cropped = F.crop(img_PIL, top, left, height, width)
-        img_cropped_GT = transforms.ToTensor()(img_PIL_cropped)
+        img_cropped_GT = transforms.ToTensor()(img_PIL_cropped).to(device)
         self.assertTrue(torch.equal(img_tensor, img_tensor_clone))
         self.assertTrue(torch.equal(img_cropped, (img_cropped_GT * 255).to(torch.uint8)),
                         "functional_tensor crop not working")
@@ -203,17 +219,15 @@ def test_center_crop(self):
         img_tensor_clone = img_tensor.clone()
         cropped_tensor = F_t.center_crop(img_tensor, [10, 10])
         cropped_pil_image = F.center_crop(transforms.ToPILImage()(img_tensor), [10, 10])
-        cropped_pil_tensor = (transforms.ToTensor()(cropped_pil_image) * 255).to(torch.uint8)
+        cropped_pil_tensor = (transforms.ToTensor()(cropped_pil_image) * 255).to(dtype=torch.uint8)
         self.assertTrue(torch.equal(cropped_tensor, cropped_pil_tensor))
         self.assertTrue(torch.equal(img_tensor, img_tensor_clone))
         # scriptable function test
         cropped_script = script_center_crop(img_tensor, [10, 10])
         self.assertTrue(torch.equal(cropped_script, cropped_tensor))
 
     def test_five_crop(self):
-        script_five_crop = torch.jit.script(F_t.five_crop)
         img_tensor = torch.randint(0, 255, (1, 32, 32), dtype=torch.uint8)
-        img_tensor_clone = img_tensor.clone()
         cropped_tensor = F_t.five_crop(img_tensor, [10, 10])
         cropped_pil_image = F.five_crop(transforms.ToPILImage()(img_tensor), [10, 10])
         self.assertTrue(torch.equal(cropped_tensor[0],
@@ -226,11 +240,6 @@ def test_five_crop(self):
                                     (transforms.ToTensor()(cropped_pil_image[3]) * 255).to(torch.uint8)))
         self.assertTrue(torch.equal(cropped_tensor[4],
                                     (transforms.ToTensor()(cropped_pil_image[4]) * 255).to(torch.uint8)))
-        self.assertTrue(torch.equal(img_tensor, img_tensor_clone))
-        # scriptable function test
-        cropped_script = script_five_crop(img_tensor, [10, 10])
-        for cropped_script_img, cropped_tensor_img in zip(cropped_script, cropped_tensor):
-            self.assertTrue(torch.equal(cropped_script_img, cropped_tensor_img))
 
     def test_ten_crop(self):
         script_ten_crop = torch.jit.script(F_t.ten_crop)
@@ -264,9 +273,10 @@ def test_ten_crop(self):
         for cropped_script_img, cropped_tensor_img in zip(cropped_script, cropped_tensor):
             self.assertTrue(torch.equal(cropped_script_img, cropped_tensor_img))
 
-    def test_pad(self):
+    @run_on_cpu_and_cuda
+    def test_pad(self, device):
         script_fn = torch.jit.script(F_t.pad)
-        tensor, pil_img = self._create_data(7, 8)
+        tensor, pil_img = self._create_data(7, 8, device=device)
 
         for dt in [None, torch.float32, torch.float64]:
             if dt is not None:
@@ -302,9 +312,10 @@ def test_pad(self):
         with self.assertRaises(ValueError, msg="Padding can not be negative for symmetric padding_mode"):
             F_t.pad(tensor, (-2, -3), padding_mode="symmetric")
 
-    def test_adjust_gamma(self):
-        script_fn = torch.jit.script(F_t.adjust_gamma)
-        tensor, pil_img = self._create_data(26, 36)
+    @run_on_cpu_and_cuda
+    def test_adjust_gamma(self, device):
+        script_fn = torch.jit.script(F.adjust_gamma)
+        tensor, pil_img = self._create_data(26, 36, device=device)
 
         for dt in [torch.float64, torch.float32, None]:
 
@@ -315,8 +326,8 @@ def test_adjust_gamma(self):
             gains = [0.7, 1.0, 1.3]
             for gamma, gain in zip(gammas, gains):
 
-                adjusted_tensor = F_t.adjust_gamma(tensor, gamma, gain)
-                adjusted_pil = F_pil.adjust_gamma(pil_img, gamma, gain)
+                adjusted_tensor = F.adjust_gamma(tensor, gamma, gain)
+                adjusted_pil = F.adjust_gamma(pil_img, gamma, gain)
                 scripted_result = script_fn(tensor, gamma, gain)
                 self.assertEqual(adjusted_tensor.dtype, scripted_result.dtype)
                 self.assertEqual(adjusted_tensor.size()[1:], adjusted_pil.size[::-1])
@@ -327,11 +338,12 @@ def test_adjust_gamma(self):
 
                 self.compareTensorToPIL(rbg_tensor, adjusted_pil)
 
-                self.assertTrue(adjusted_tensor.equal(scripted_result))
+                self.assertTrue(adjusted_tensor.allclose(scripted_result))
 
-    def test_resize(self):
+    @run_on_cpu_and_cuda
+    def test_resize(self, device):
         script_fn = torch.jit.script(F_t.resize)
-        tensor, pil_img = self._create_data(26, 36)
+        tensor, pil_img = self._create_data(26, 36, device=device)
 
         for dt in [None, torch.float32, torch.float64]:
             if dt is not None:
@@ -367,28 +379,30 @@ def test_resize(self):
                     resize_result = script_fn(tensor, size=script_size, interpolation=interpolation)
                     self.assertTrue(resized_tensor.equal(resize_result), msg="{}, {}".format(size, interpolation))
 
-    def test_resized_crop(self):
+    @run_on_cpu_and_cuda
+    def test_resized_crop(self, device):
         # test values of F.resized_crop in several cases:
         # 1) resize to the same size, crop to the same size => should be identity
-        tensor, _ = self._create_data(26, 36)
+        tensor, _ = self._create_data(26, 36, device=device)
         for i in [0, 2, 3]:
             out_tensor = F.resized_crop(tensor, top=0, left=0, height=26, width=36, size=[26, 36], interpolation=i)
             self.assertTrue(tensor.equal(out_tensor), msg="{} vs {}".format(out_tensor[0, :5, :5], tensor[0, :5, :5]))
 
         # 2) resize by half and crop a TL corner
-        tensor, _ = self._create_data(26, 36)
+        tensor, _ = self._create_data(26, 36, device=device)
         out_tensor = F.resized_crop(tensor, top=0, left=0, height=20, width=30, size=[10, 15], interpolation=0)
         expected_out_tensor = tensor[:, :20:2, :30:2]
         self.assertTrue(
             expected_out_tensor.equal(out_tensor),
             msg="{} vs {}".format(expected_out_tensor[0, :10, :10], out_tensor[0, :10, :10])
         )
 
-    def test_affine(self):
+    @run_on_cpu_and_cuda
+    def test_affine(self, device):
         # Tests on square and rectangular images
         scripted_affine = torch.jit.script(F.affine)
 
-        for tensor, pil_img in [self._create_data(26, 26), self._create_data(32, 26)]:
+        for tensor, pil_img in [self._create_data(26, 26, device=device), self._create_data(32, 26, device=device)]:
 
             # 1) identity map
             out_tensor = F.affine(tensor, angle=0, translate=[0, 0], scale=1.0, shear=[0.0, 0.0], resample=0)
@@ -412,8 +426,16 @@ def test_affine(self):
                     (180, torch.rot90(tensor, k=2, dims=(-1, -2))),
                 ]
                 for a, true_tensor in test_configs:
+
+                    out_pil_img = F.affine(
+                        pil_img, angle=a, translate=[0, 0], scale=1.0, shear=[0.0, 0.0], resample=0
+                    )
+                    out_pil_tensor = torch.from_numpy(np.array(out_pil_img).transpose((2, 0, 1))).to(device)
+
                     for fn in [F.affine, scripted_affine]:
-                        out_tensor = fn(tensor, angle=a, translate=[0, 0], scale=1.0, shear=[0.0, 0.0], resample=0)
+                        out_tensor = fn(
+                            tensor, angle=a, translate=[0, 0], scale=1.0, shear=[0.0, 0.0], resample=0
+                        )
                         if true_tensor is not None:
                             self.assertTrue(
                                 true_tensor.equal(out_tensor),
@@ -422,11 +444,6 @@ def test_affine(self):
                         else:
                             true_tensor = out_tensor
 
-                        out_pil_img = F.affine(
-                            pil_img, angle=a, translate=[0, 0], scale=1.0, shear=[0.0, 0.0], resample=0
-                        )
-                        out_pil_tensor = torch.from_numpy(np.array(out_pil_img).transpose((2, 0, 1)))
-
                         num_diff_pixels = (true_tensor != out_pil_tensor).sum().item() / 3.0
                         ratio_diff_pixels = num_diff_pixels / true_tensor.shape[-1] / true_tensor.shape[-2]
                         # Tolerance : less than 6% of different pixels
@@ -442,12 +459,16 @@ def test_affine(self):
                     90, 45, 15, -30, -60, -120
                 ]
                 for a in test_configs:
+
+                    out_pil_img = F.affine(
+                        pil_img, angle=a, translate=[0, 0], scale=1.0, shear=[0.0, 0.0], resample=0
+                    )
+                    out_pil_tensor = torch.from_numpy(np.array(out_pil_img).transpose((2, 0, 1)))
+
                     for fn in [F.affine, scripted_affine]:
-                        out_tensor = fn(tensor, angle=a, translate=[0, 0], scale=1.0, shear=[0.0, 0.0], resample=0)
-                        out_pil_img = F.affine(
-                            pil_img, angle=a, translate=[0, 0], scale=1.0, shear=[0.0, 0.0], resample=0
-                        )
-                        out_pil_tensor = torch.from_numpy(np.array(out_pil_img).transpose((2, 0, 1)))
+                        out_tensor = fn(
+                            tensor, angle=a, translate=[0, 0], scale=1.0, shear=[0.0, 0.0], resample=0
+                        ).cpu()
 
                         num_diff_pixels = (out_tensor != out_pil_tensor).sum().item() / 3.0
                         ratio_diff_pixels = num_diff_pixels / out_tensor.shape[-1] / out_tensor.shape[-2]
@@ -465,9 +486,12 @@ def test_affine(self):
                 [10, 12], (-12, -13)
             ]
             for t in test_configs:
+
+                out_pil_img = F.affine(pil_img, angle=0, translate=t, scale=1.0, shear=[0.0, 0.0], resample=0)
+
                 for fn in [F.affine, scripted_affine]:
                     out_tensor = fn(tensor, angle=0, translate=t, scale=1.0, shear=[0.0, 0.0], resample=0)
-                    out_pil_img = F.affine(pil_img, angle=0, translate=t, scale=1.0, shear=[0.0, 0.0], resample=0)
+
                     self.compareTensorToPIL(out_tensor, out_pil_img)
 
             # 3) Test rotation + translation + scale + share
@@ -489,23 +513,25 @@ def test_affine(self):
                     out_pil_tensor = torch.from_numpy(np.array(out_pil_img).transpose((2, 0, 1)))
 
                     for fn in [F.affine, scripted_affine]:
-                        out_tensor = fn(tensor, angle=a, translate=t, scale=s, shear=sh, resample=r)
+                        out_tensor = fn(tensor, angle=a, translate=t, scale=s, shear=sh, resample=r).cpu()
                         num_diff_pixels = (out_tensor != out_pil_tensor).sum().item() / 3.0
                         ratio_diff_pixels = num_diff_pixels / out_tensor.shape[-1] / out_tensor.shape[-2]
-                        # Tolerance : less than 5% of different pixels
+                        # Tolerance : less than 5% (cpu), 6% (cuda) of different pixels
+                        tol = 0.06 if device == "cuda" else 0.05
                         self.assertLess(
                             ratio_diff_pixels,
-                            0.05,
+                            tol,
                             msg="{}: {}\n{} vs \n{}".format(
                                 (r, a, t, s, sh), ratio_diff_pixels, out_tensor[0, :7, :7], out_pil_tensor[0, :7, :7]
                             )
                         )
 
-    def test_rotate(self):
+    @run_on_cpu_and_cuda
+    def test_rotate(self, device):
         # Tests on square image
         scripted_rotate = torch.jit.script(F.rotate)
 
-        for tensor, pil_img in [self._create_data(26, 26), self._create_data(32, 26)]:
+        for tensor, pil_img in [self._create_data(26, 26, device=device), self._create_data(32, 26, device=device)]:
 
             img_size = pil_img.size
             centers = [
@@ -522,7 +548,7 @@ def test_rotate(self):
                             out_pil_img = F.rotate(pil_img, angle=a, resample=r, expand=e, center=c)
                             out_pil_tensor = torch.from_numpy(np.array(out_pil_img).transpose((2, 0, 1)))
                             for fn in [F.rotate, scripted_rotate]:
-                                out_tensor = fn(tensor, angle=a, resample=r, expand=e, center=c)
+                                out_tensor = fn(tensor, angle=a, resample=r, expand=e, center=c).cpu()
 
                                 self.assertEqual(
                                     out_tensor.shape,
@@ -545,11 +571,12 @@ def test_rotate(self):
                                     )
                                 )
 
-    def test_perspective(self):
+    @run_on_cpu_and_cuda
+    def test_perspective(self, device):
 
         from torchvision.transforms import RandomPerspective
 
-        for tensor, pil_img in [self._create_data(26, 34), self._create_data(26, 26)]:
+        for tensor, pil_img in [self._create_data(26, 34, device=device), self._create_data(26, 26, device=device)]:
 
             scripted_tranform = torch.jit.script(F.perspective)
 
@@ -569,7 +596,7 @@ def test_perspective(self):
                     out_pil_tensor = torch.from_numpy(np.array(out_pil_img).transpose((2, 0, 1)))
 
                     for fn in [F.perspective, scripted_tranform]:
-                        out_tensor = fn(tensor, startpoints=spoints, endpoints=epoints, interpolation=r)
+                        out_tensor = fn(tensor, startpoints=spoints, endpoints=epoints, interpolation=r).cpu()
 
                         num_diff_pixels = (out_tensor != out_pil_tensor).sum().item() / 3.0
                         ratio_diff_pixels = num_diff_pixels / out_tensor.shape[-1] / out_tensor.shape[-2]

torchvision/transforms/functional.py

@@ -223,9 +223,11 @@ def to_pil_image(pic, mode=None):
             pic = np.expand_dims(pic, 2)
 
     npimg = pic
-    if isinstance(pic, torch.FloatTensor) and mode != 'F':
+    if pic.is_floating_point() and mode != 'F':
         pic = pic.mul(255).byte()
     if isinstance(pic, torch.Tensor):
+        if pic.device != torch.device("cpu"):
+            pic = pic.cpu()
         npimg = np.transpose(pic.numpy(), (1, 2, 0))
 
     if not isinstance(npimg, np.ndarray):

torchvision/transforms/functional_tensor.py

@@ -659,6 +659,9 @@ def _apply_grid_transform(img: Tensor, grid: Tensor, mode: str) -> Tensor:
         need_cast = True
         img = img.to(torch.float32)
 
+    if grid.device.type != img.device.type:
+        grid = grid.to(img)
+
     img = grid_sample(img, grid, mode=mode, padding_mode="zeros", align_corners=False)
 
     if need_squeeze: