diff --git a/test/test_functional_tensor.py b/test/test_functional_tensor.py index 68359bc0437..5994cda792b 100644 --- a/test/test_functional_tensor.py +++ b/test/test_functional_tensor.py @@ -16,9 +16,12 @@ class Tester(TransformsTester): - def _test_vflip(self, device): + def setUp(self): + self.device = "cpu" + + def test_vflip(self): script_vflip = torch.jit.script(F_t.vflip) - img_tensor = torch.randn(3, 16, 16, device=device) + img_tensor = torch.randn(3, 16, 16, device=self.device) img_tensor_clone = img_tensor.clone() vflipped_img = F_t.vflip(img_tensor) vflipped_img_again = F_t.vflip(vflipped_img) @@ -29,16 +32,9 @@ def _test_vflip(self, device): vflipped_img_script = script_vflip(img_tensor) self.assertTrue(torch.equal(vflipped_img, vflipped_img_script)) - def test_vflip_cpu(self): - self._test_vflip("cpu") - - @unittest.skipIf(not torch.cuda.is_available(), reason="Skip if no CUDA device") - def test_vflip_cuda(self): - self._test_vflip("cuda") - - def _test_hflip(self, device): + def test_hflip(self): script_hflip = torch.jit.script(F_t.hflip) - img_tensor = torch.randn(3, 16, 16, device=device) + img_tensor = torch.randn(3, 16, 16, device=self.device) img_tensor_clone = img_tensor.clone() hflipped_img = F_t.hflip(img_tensor) hflipped_img_again = F_t.hflip(hflipped_img) @@ -49,17 +45,10 @@ def _test_hflip(self, device): hflipped_img_script = script_hflip(img_tensor) self.assertTrue(torch.equal(hflipped_img, hflipped_img_script)) - def test_hflip_cpu(self): - self._test_hflip("cpu") - - @unittest.skipIf(not torch.cuda.is_available(), reason="Skip if no CUDA device") - def test_hflip_cuda(self): - self._test_hflip("cuda") - - def _test_crop(self, device): + def test_crop(self): script_crop = torch.jit.script(F.crop) - img_tensor, pil_img = self._create_data(16, 18, device=device) + img_tensor, pil_img = self._create_data(16, 18, device=self.device) test_configs = [ (1, 2, 4, 5), # crop inside top-left corner @@ -77,13 +66,6 @@ def _test_crop(self, device): img_tensor_cropped = script_crop(img_tensor, top, left, height, width) self.compareTensorToPIL(img_tensor_cropped, pil_img_cropped) - def test_crop_cpu(self): - self._test_crop("cpu") - - @unittest.skipIf(not torch.cuda.is_available(), reason="Skip if no CUDA device") - def test_crop_cuda(self): - self._test_crop("cuda") - def test_hsv2rgb(self): shape = (3, 100, 150) for _ in range(20): @@ -129,7 +111,7 @@ def test_rgb2hsv(self): self.assertLess(max_diff, 1e-5) - def _test_adjustments(self, device): + def test_adjustments(self): script_adjust_brightness = torch.jit.script(F_t.adjust_brightness) script_adjust_contrast = torch.jit.script(F_t.adjust_contrast) script_adjust_saturation = torch.jit.script(F_t.adjust_saturation) @@ -144,9 +126,9 @@ def _test_adjustments(self, device): shape = (channels, dims[0], dims[1]) if torch.randint(0, 2, (1,)) == 0: - img = torch.rand(*shape, dtype=torch.float, device=device) + img = torch.rand(*shape, dtype=torch.float, device=self.device) else: - img = torch.randint(0, 256, shape, dtype=torch.uint8, device=device) + img = torch.randint(0, 256, shape, dtype=torch.uint8, device=self.device) factor = 3 * torch.rand(1).item() img_clone = img.clone() @@ -187,17 +169,10 @@ def _test_adjustments(self, device): scripted_fn = torch.jit.script(f) scripted_fn(img) - def test_adjustments(self): - self._test_adjustments("cpu") - - @unittest.skipIf(not torch.cuda.is_available(), reason="Skip if no CUDA device") - def test_adjustments_cuda(self): - self._test_adjustments("cuda") - - def _test_rgb_to_grayscale(self, device): + def test_rgb_to_grayscale(self): script_rgb_to_grayscale = torch.jit.script(F.rgb_to_grayscale) - img_tensor, pil_img = self._create_data(32, 34, device=device) + img_tensor, pil_img = self._create_data(32, 34, device=self.device) for num_output_channels in (3, 1): gray_pil_image = F.rgb_to_grayscale(pil_img, num_output_channels=num_output_channels) @@ -211,17 +186,10 @@ def _test_rgb_to_grayscale(self, device): s_gray_tensor = script_rgb_to_grayscale(img_tensor, num_output_channels=num_output_channels) self.assertTrue(s_gray_tensor.equal(gray_tensor)) - def test_rgb_to_grayscale(self): - self._test_rgb_to_grayscale("cpu") - - @unittest.skipIf(not torch.cuda.is_available(), reason="Skip if no CUDA device") - def test_rgb_to_grayscale_cuda(self): - self._test_rgb_to_grayscale("cuda") - - def _test_center_crop(self, device): + def test_center_crop(self): script_center_crop = torch.jit.script(F.center_crop) - img_tensor, pil_img = self._create_data(32, 34, device=device) + img_tensor, pil_img = self._create_data(32, 34, device=self.device) cropped_pil_image = F.center_crop(pil_img, [10, 11]) @@ -231,17 +199,10 @@ def _test_center_crop(self, device): cropped_tensor = script_center_crop(img_tensor, [10, 11]) self.compareTensorToPIL(cropped_tensor, cropped_pil_image) - def test_center_crop(self): - self._test_center_crop("cpu") - - @unittest.skipIf(not torch.cuda.is_available(), reason="Skip if no CUDA device") - def test_center_crop_cuda(self): - self._test_center_crop("cuda") - - def _test_five_crop(self, device): + def test_five_crop(self): script_five_crop = torch.jit.script(F.five_crop) - img_tensor, pil_img = self._create_data(32, 34, device=device) + img_tensor, pil_img = self._create_data(32, 34, device=self.device) cropped_pil_images = F.five_crop(pil_img, [10, 11]) @@ -253,17 +214,10 @@ def _test_five_crop(self, device): for i in range(5): self.compareTensorToPIL(cropped_tensors[i], cropped_pil_images[i]) - def test_five_crop(self): - self._test_five_crop("cpu") - - @unittest.skipIf(not torch.cuda.is_available(), reason="Skip if no CUDA device") - def test_five_crop_cuda(self): - self._test_five_crop("cuda") - - def _test_ten_crop(self, device): + def test_ten_crop(self): script_ten_crop = torch.jit.script(F.ten_crop) - img_tensor, pil_img = self._create_data(32, 34, device=device) + img_tensor, pil_img = self._create_data(32, 34, device=self.device) cropped_pil_images = F.ten_crop(pil_img, [10, 11]) @@ -275,16 +229,9 @@ def _test_ten_crop(self, device): for i in range(10): self.compareTensorToPIL(cropped_tensors[i], cropped_pil_images[i]) - def test_ten_crop(self): - self._test_ten_crop("cpu") - - @unittest.skipIf(not torch.cuda.is_available(), reason="Skip if no CUDA device") - def test_ten_crop_cuda(self): - self._test_ten_crop("cuda") - - def _test_pad(self, device): + def test_pad(self): script_fn = torch.jit.script(F_t.pad) - tensor, pil_img = self._create_data(7, 8, device=device) + tensor, pil_img = self._create_data(7, 8, device=self.device) for dt in [None, torch.float32, torch.float64]: if dt is not None: @@ -320,16 +267,9 @@ def _test_pad(self, device): 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_pad_cpu(self): - self._test_pad("cpu") - - @unittest.skipIf(not torch.cuda.is_available(), reason="Skip if no CUDA device") - def test_pad_cuda(self): - self._test_pad("cuda") - - def _test_adjust_gamma(self, device): + def test_adjust_gamma(self): script_fn = torch.jit.script(F.adjust_gamma) - tensor, pil_img = self._create_data(26, 36, device=device) + tensor, pil_img = self._create_data(26, 36, device=self.device) for dt in [torch.float64, torch.float32, None]: @@ -354,16 +294,9 @@ def _test_adjust_gamma(self, device): self.assertTrue(adjusted_tensor.allclose(scripted_result)) - def test_adjust_gamma_cpu(self): - self._test_adjust_gamma("cpu") - - @unittest.skipIf(not torch.cuda.is_available(), reason="Skip if no CUDA device") - def test_adjust_gamma_cuda(self): - self._test_adjust_gamma("cuda") - - def _test_resize(self, device): + def test_resize(self): script_fn = torch.jit.script(F_t.resize) - tensor, pil_img = self._create_data(26, 36, device=device) + tensor, pil_img = self._create_data(26, 36, device=self.device) for dt in [None, torch.float32, torch.float64]: if dt is not None: @@ -399,23 +332,16 @@ def _test_resize(self, device): resize_result = script_fn(tensor, size=script_size, interpolation=interpolation) self.assertTrue(resized_tensor.equal(resize_result), msg="{}, {}".format(size, interpolation)) - def test_resize_cpu(self): - self._test_resize("cpu") - - @unittest.skipIf(not torch.cuda.is_available(), reason="Skip if no CUDA device") - def test_resize_cuda(self): - self._test_resize("cuda") - - def _test_resized_crop(self, device): + def test_resized_crop(self): # 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, device=device) + tensor, _ = self._create_data(26, 36, device=self.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, device=device) + tensor, _ = self._create_data(26, 36, device=self.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( @@ -423,18 +349,12 @@ def _test_resized_crop(self, device): msg="{} vs {}".format(expected_out_tensor[0, :10, :10], out_tensor[0, :10, :10]) ) - def test_resized_crop_cpu(self): - self._test_resized_crop("cpu") - - @unittest.skipIf(not torch.cuda.is_available(), reason="Skip if no CUDA device") - def test_resized_crop_cuda(self): - self._test_resized_crop("cuda") - - def _test_affine(self, device): + def test_affine(self): # Tests on square and rectangular images scripted_affine = torch.jit.script(F.affine) - for tensor, pil_img in [self._create_data(26, 26, device=device), self._create_data(32, 26, device=device)]: + data = [self._create_data(26, 26, device=self.device), self._create_data(32, 26, device=self.device)] + for tensor, pil_img in data: # 1) identity map out_tensor = F.affine(tensor, angle=0, translate=[0, 0], scale=1.0, shear=[0.0, 0.0], resample=0) @@ -462,7 +382,7 @@ def _test_affine(self, device): 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) + out_pil_tensor = torch.from_numpy(np.array(out_pil_img).transpose((2, 0, 1))).to(self.device) for fn in [F.affine, scripted_affine]: out_tensor = fn( @@ -549,7 +469,7 @@ def _test_affine(self, device): 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% (cpu), 6% (cuda) of different pixels - tol = 0.06 if device == "cuda" else 0.05 + tol = 0.06 if self.device == "cuda" else 0.05 self.assertLess( ratio_diff_pixels, tol, @@ -558,18 +478,12 @@ def _test_affine(self, device): ) ) - def test_affine_cpu(self): - self._test_affine("cpu") - - @unittest.skipIf(not torch.cuda.is_available(), reason="Skip if no CUDA device") - def test_affine_cuda(self): - self._test_affine("cuda") - - def _test_rotate(self, device): + def test_rotate(self): # Tests on square image scripted_rotate = torch.jit.script(F.rotate) - for tensor, pil_img in [self._create_data(26, 26, device=device), self._create_data(32, 26, device=device)]: + data = [self._create_data(26, 26, device=self.device), self._create_data(32, 26, device=self.device)] + for tensor, pil_img in data: img_size = pil_img.size centers = [ @@ -609,18 +523,12 @@ def _test_rotate(self, device): ) ) - def test_rotate_cpu(self): - self._test_rotate("cpu") - - @unittest.skipIf(not torch.cuda.is_available(), reason="Skip if no CUDA device") - def test_rotate_cuda(self): - self._test_rotate("cuda") - - def _test_perspective(self, device): + def test_perspective(self): from torchvision.transforms import RandomPerspective - for tensor, pil_img in [self._create_data(26, 34, device=device), self._create_data(26, 26, device=device)]: + data = [self._create_data(26, 34, device=self.device), self._create_data(26, 26, device=self.device)] + for tensor, pil_img in data: scripted_tranform = torch.jit.script(F.perspective) @@ -656,12 +564,12 @@ def _test_perspective(self, device): ) ) - def test_perspective_cpu(self): - self._test_perspective("cpu") - @unittest.skipIf(not torch.cuda.is_available(), reason="Skip if no CUDA device") - def test_perspective_cuda(self): - self._test_perspective("cuda") +@unittest.skipIf(not torch.cuda.is_available(), reason="Skip if no CUDA device") +class CUDATester(Tester): + + def setUp(self): + self.device = "cuda" if __name__ == '__main__': diff --git a/test/test_transforms_tensor.py b/test/test_transforms_tensor.py index 245a27954f9..de24ec4912e 100644 --- a/test/test_transforms_tensor.py +++ b/test/test_transforms_tensor.py @@ -13,10 +13,13 @@ class Tester(TransformsTester): + def setUp(self): + self.device = "cpu" + def _test_functional_op(self, func, fn_kwargs): if fn_kwargs is None: fn_kwargs = {} - tensor, pil_img = self._create_data(height=10, width=10) + tensor, pil_img = self._create_data(height=10, width=10, device=self.device) transformed_tensor = getattr(F, func)(tensor, **fn_kwargs) transformed_pil_img = getattr(F, func)(pil_img, **fn_kwargs) self.compareTensorToPIL(transformed_tensor, transformed_pil_img) @@ -25,7 +28,7 @@ def _test_class_op(self, method, meth_kwargs=None, test_exact_match=True, **matc if meth_kwargs is None: meth_kwargs = {} - tensor, pil_img = self._create_data(height=10, width=10) + tensor, pil_img = self._create_data(height=10, width=10, device=self.device) # test for class interface f = getattr(T, method)(**meth_kwargs) scripted_fn = torch.jit.script(f) @@ -58,14 +61,14 @@ def test_adjustments(self): fns = ['adjust_brightness', 'adjust_contrast', 'adjust_saturation'] for _ in range(20): factor = 3 * torch.rand(1).item() - tensor, _ = self._create_data() + tensor, _ = self._create_data(device=self.device) pil_img = T.ToPILImage()(tensor) for func in fns: adjusted_tensor = getattr(F, func)(tensor, factor) adjusted_pil_img = getattr(F, func)(pil_img, factor) - adjusted_pil_tensor = T.ToTensor()(adjusted_pil_img) + adjusted_pil_tensor = T.ToTensor()(adjusted_pil_img).to(self.device) scripted_fn = torch.jit.script(getattr(F, func)) adjusted_tensor_script = scripted_fn(tensor, factor) @@ -136,7 +139,7 @@ def test_center_crop(self): self._test_op( "center_crop", "CenterCrop", fn_kwargs=fn_kwargs, meth_kwargs=meth_kwargs ) - tensor = torch.randint(0, 255, (3, 10, 10), dtype=torch.uint8) + tensor = torch.randint(0, 255, (3, 10, 10), dtype=torch.uint8, device=self.device) # Test torchscript of transforms.CenterCrop with size as int f = T.CenterCrop(size=5) scripted_fn = torch.jit.script(f) @@ -157,7 +160,7 @@ def _test_op_list_output(self, func, method, out_length, fn_kwargs=None, meth_kw fn_kwargs = {} if meth_kwargs is None: meth_kwargs = {} - tensor, pil_img = self._create_data(height=20, width=20) + tensor, pil_img = self._create_data(height=20, width=20, device=self.device) transformed_t_list = getattr(F, func)(tensor, **fn_kwargs) transformed_p_list = getattr(F, func)(pil_img, **fn_kwargs) self.assertEqual(len(transformed_t_list), len(transformed_p_list)) @@ -216,7 +219,7 @@ def test_ten_crop(self): ) def test_resize(self): - tensor, _ = self._create_data(height=34, width=36) + tensor, _ = self._create_data(height=34, width=36, device=self.device) script_fn = torch.jit.script(F.resize) for dt in [None, torch.float32, torch.float64]: @@ -243,7 +246,7 @@ def test_resize(self): self.assertTrue(s_resized_tensor.equal(resized_tensor)) def test_resized_crop(self): - tensor = torch.randint(0, 255, size=(3, 44, 56), dtype=torch.uint8) + tensor = torch.randint(0, 255, size=(3, 44, 56), dtype=torch.uint8, device=self.device) for scale in [(0.7, 1.2), [0.7, 1.2]]: for ratio in [(0.75, 1.333), [0.75, 1.333]]: @@ -261,7 +264,7 @@ def test_resized_crop(self): self.assertTrue(out1.equal(out2)) def test_random_affine(self): - tensor = torch.randint(0, 255, size=(3, 44, 56), dtype=torch.uint8) + tensor = torch.randint(0, 255, size=(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]]: for scale in [(0.7, 1.2), [0.7, 1.2]]: @@ -281,7 +284,7 @@ def test_random_affine(self): self.assertTrue(out1.equal(out2)) def test_random_rotate(self): - tensor = torch.randint(0, 255, size=(3, 44, 56), dtype=torch.uint8) + tensor = torch.randint(0, 255, size=(3, 44, 56), dtype=torch.uint8, device=self.device) for center in [(0, 0), [10, 10], None, (56, 44)]: for expand in [True, False]: @@ -299,7 +302,7 @@ def test_random_rotate(self): self.assertTrue(out1.equal(out2)) def test_random_perspective(self): - tensor = torch.randint(0, 255, size=(3, 44, 56), dtype=torch.uint8) + tensor = torch.randint(0, 255, size=(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]: @@ -334,5 +337,12 @@ def test_to_grayscale(self): ) +@unittest.skipIf(not torch.cuda.is_available(), reason="Skip if no CUDA device") +class CUDATester(Tester): + + def setUp(self): + self.device = "cuda" + + if __name__ == '__main__': unittest.main()