has cupy check

wyli · wyli · commit 4190a3eec530 · 2021-12-23T22:37:27.000Z
Signed-off-by: Wenqi Li &lt;wenqil@nvidia.com&gt;
diff --git a/tests/test_cast_to_type.py b/tests/test_cast_to_type.py
@@ -16,11 +16,8 @@
 from parameterized import parameterized
 
 from monai.transforms import CastToType
-from monai.utils import optional_import
 from monai.utils.type_conversion import get_equivalent_dtype
-from tests.utils import TEST_NDARRAYS
-
-cp, has_cp = optional_import("cupy")
+from tests.utils import HAS_CUPY, TEST_NDARRAYS
 
 TESTS = []
 for p in TEST_NDARRAYS:
@@ -45,7 +42,7 @@ def test_type(self, out_dtype, input_data, expected_type):
         self.assertEqual(result.dtype, get_equivalent_dtype(expected_type, type(result)))
 
     @parameterized.expand(TESTS_CUPY)
-    @unittest.skipUnless(has_cp, "Requires CuPy")
+    @unittest.skipUnless(HAS_CUPY, "Requires CuPy")
     def test_type_cupy(self, out_dtype, input_data, expected_type):
         input_data = cp.asarray(input_data)
 
diff --git a/tests/test_cast_to_typed.py b/tests/test_cast_to_typed.py
@@ -16,9 +16,7 @@
 from parameterized import parameterized
 
 from monai.transforms import CastToTyped
-from monai.utils import optional_import
-
-cp, has_cp = optional_import("cupy")
+from tests.utils import HAS_CUPY
 
 TEST_CASE_1 = [
     {"keys": ["img"], "dtype": np.float64},
@@ -58,7 +56,7 @@ def test_type(self, input_param, input_data, expected_type):
             self.assertEqual(v.dtype, expected_type[k])
 
     @parameterized.expand(TESTS_CUPY)
-    @unittest.skipUnless(has_cp, "Requires CuPy")
+    @unittest.skipUnless(HAS_CUPY, "Requires CuPy")
     def test_type_cupy(self, input_param, input_data, expected_type):
         input_data = {k: cp.asarray(v) for k, v in input_data.items()}
 
diff --git a/tests/test_cucim_dict_transform.py b/tests/test_cucim_dict_transform.py
@@ -16,10 +16,9 @@
 
 from monai.transforms import CuCIMd
 from monai.utils import optional_import, set_determinism
-from tests.utils import skip_if_no_cuda
+from tests.utils import HAS_CUPY, skip_if_no_cuda
 
 _, has_cut = optional_import("cucim.core.operations.expose.transform")
-cp, has_cp = optional_import("cupy")
 
 set_determinism(seed=0)
 
@@ -62,7 +61,7 @@
 
 
 @skip_if_no_cuda
-@unittest.skipUnless(has_cp, "CuPy is required.")
+@unittest.skipUnless(HAS_CUPY, "CuPy is required.")
 @unittest.skipUnless(has_cut, "cuCIM transforms are required.")
 class TestCuCIMDict(unittest.TestCase):
     @parameterized.expand(
diff --git a/tests/test_cucim_transform.py b/tests/test_cucim_transform.py
@@ -16,10 +16,9 @@
 
 from monai.transforms import CuCIM
 from monai.utils import optional_import, set_determinism
-from tests.utils import skip_if_no_cuda
+from tests.utils import HAS_CUPY, skip_if_no_cuda
 
 _, has_cut = optional_import("cucim.core.operations.expose.transform")
-cp, has_cp = optional_import("cupy")
 
 set_determinism(seed=0)
 
@@ -62,7 +61,7 @@
 
 
 @skip_if_no_cuda
-@unittest.skipUnless(has_cp, "CuPy is required.")
+@unittest.skipUnless(HAS_CUPY, "CuPy is required.")
 @unittest.skipUnless(has_cut, "cuCIM transforms are required.")
 class TestCuCIM(unittest.TestCase):
     @parameterized.expand(
diff --git a/tests/test_nvtx_decorator.py b/tests/test_nvtx_decorator.py
@@ -31,9 +31,9 @@
     ToTensorD,
 )
 from monai.utils import Range, optional_import
+from tests.utils import HAS_CUPY
 
 _, has_nvtx = optional_import("torch._C._nvtx", descriptor="NVTX is not installed. Are you sure you have a CUDA build?")
-_, has_cp = optional_import("cupy")
 _, has_tvt = optional_import("torchvision.transforms")
 _, has_cut = optional_import("cucim.core.operations.expose.transform")
 
@@ -103,7 +103,7 @@ def test_tranform_dict(self, input):
         np.testing.assert_equal(output.numpy(), output3.numpy())
 
     @parameterized.expand([TEST_CASE_WRAPPER])
-    @unittest.skipUnless(has_cp, "Requires CuPy.")
+    @unittest.skipUnless(HAS_CUPY, "Requires CuPy.")
     @unittest.skipUnless(has_cut, "Requires cuCIM transforms.")
     @unittest.skipUnless(has_tvt, "Requires torchvision transforms.")
     def test_wrapper_tranforms(self, input):
diff --git a/tests/test_rand_cucim_dict_transform.py b/tests/test_rand_cucim_dict_transform.py
@@ -16,10 +16,9 @@
 
 from monai.transforms import RandCuCIMd
 from monai.utils import optional_import, set_determinism
-from tests.utils import skip_if_no_cuda
+from tests.utils import HAS_CUPY, skip_if_no_cuda
 
 _, has_cut = optional_import("cucim.core.operations.expose.transform")
-cp, has_cp = optional_import("cupy")
 
 set_determinism(seed=0)
 
@@ -74,7 +73,7 @@
 
 
 @skip_if_no_cuda
-@unittest.skipUnless(has_cp, "CuPy is required.")
+@unittest.skipUnless(HAS_CUPY, "CuPy is required.")
 @unittest.skipUnless(has_cut, "cuCIM transforms are required.")
 class TestRandCuCIMDict(unittest.TestCase):
     @parameterized.expand(
diff --git a/tests/test_rand_cucim_transform.py b/tests/test_rand_cucim_transform.py
@@ -16,10 +16,9 @@
 
 from monai.transforms import RandCuCIM
 from monai.utils import optional_import, set_determinism
-from tests.utils import skip_if_no_cuda
+from tests.utils import HAS_CUPY, skip_if_no_cuda
 
 _, has_cut = optional_import("cucim.core.operations.expose.transform")
-cp, has_cp = optional_import("cupy")
 
 set_determinism(seed=0)
 
@@ -74,7 +73,7 @@
 
 
 @skip_if_no_cuda
-@unittest.skipUnless(has_cp, "CuPy is required.")
+@unittest.skipUnless(HAS_CUPY, "CuPy is required.")
 @unittest.skipUnless(has_cut, "cuCIM transforms are required.")
 class TestRandCuCIM(unittest.TestCase):
     @parameterized.expand(
diff --git a/tests/test_to_cupy.py b/tests/test_to_cupy.py
@@ -16,13 +16,10 @@
 import torch
 
 from monai.transforms import ToCupy
-from monai.utils import optional_import
-from tests.utils import skip_if_no_cuda
+from tests.utils import HAS_CUPY, skip_if_no_cuda
 
-cp, has_cp = optional_import("cupy")
 
-
-@skipUnless(has_cp, "CuPy is required.")
+@skipUnless(HAS_CUPY, "CuPy is required.")
 class TestToCupy(unittest.TestCase):
     def test_cupy_input(self):
         test_data = cp.array([[1, 2], [3, 4]], dtype=cp.float32)
diff --git a/tests/test_to_cupyd.py b/tests/test_to_cupyd.py
@@ -16,14 +16,11 @@
 import torch
 
 from monai.transforms import ToCupyd
-from monai.utils import optional_import
-from tests.utils import skip_if_no_cuda
-
-cp, has_cp = optional_import("cupy")
+from tests.utils import HAS_CUPY, skip_if_no_cuda
 
 
+@skipUnless(HAS_CUPY, "requires CUPY")
 class TestToCupyd(unittest.TestCase):
-    @skipUnless(has_cp, "CuPy is required.")
     def test_cupy_input(self):
         test_data = cp.array([[1, 2], [3, 4]])
         test_data = cp.rot90(test_data)
@@ -33,7 +30,6 @@ def test_cupy_input(self):
         self.assertTrue(result.flags["C_CONTIGUOUS"])
         cp.testing.assert_allclose(result, test_data)
 
-    @skipUnless(has_cp, "CuPy is required.")
     def test_numpy_input(self):
         test_data = np.array([[1, 2], [3, 4]])
         test_data = np.rot90(test_data)
@@ -43,7 +39,6 @@ def test_numpy_input(self):
         self.assertTrue(result.flags["C_CONTIGUOUS"])
         cp.testing.assert_allclose(result, test_data)
 
-    @skipUnless(has_cp, "CuPy is required.")
     def test_tensor_input(self):
         test_data = torch.tensor([[1, 2], [3, 4]])
         test_data = test_data.rot90()
@@ -53,7 +48,6 @@ def test_tensor_input(self):
         self.assertTrue(result.flags["C_CONTIGUOUS"])
         cp.testing.assert_allclose(result, test_data.numpy())
 
-    @skipUnless(has_cp, "CuPy is required.")
     @skip_if_no_cuda
     def test_tensor_cuda_input(self):
         test_data = torch.tensor([[1, 2], [3, 4]]).cuda()
@@ -64,7 +58,6 @@ def test_tensor_cuda_input(self):
         self.assertTrue(result.flags["C_CONTIGUOUS"])
         cp.testing.assert_allclose(result, test_data.cpu().numpy())
 
-    @skipUnless(has_cp, "CuPy is required.")
     def test_list_tuple(self):
         test_data = [[1, 2], [3, 4]]
         result = ToCupyd(keys="img", wrap_sequence=True)({"img": test_data})["img"]
diff --git a/tests/test_to_numpy.py b/tests/test_to_numpy.py
@@ -16,14 +16,11 @@
 import torch
 
 from monai.transforms import ToNumpy
-from monai.utils import optional_import
-from tests.utils import assert_allclose, skip_if_no_cuda
-
-cp, has_cp = optional_import("cupy")
+from tests.utils import HAS_CUPY, assert_allclose, skip_if_no_cuda
 
 
 class TestToNumpy(unittest.TestCase):
-    @skipUnless(has_cp, "CuPy is required.")
+    @skipUnless(HAS_CUPY, "CuPy is required.")
     def test_cupy_input(self):
         test_data = cp.array([[1, 2], [3, 4]])
         test_data = cp.rot90(test_data)
diff --git a/tests/test_to_numpyd.py b/tests/test_to_numpyd.py
@@ -16,14 +16,11 @@
 import torch
 
 from monai.transforms import ToNumpyd
-from monai.utils import optional_import
-from tests.utils import assert_allclose, skip_if_no_cuda
-
-cp, has_cp = optional_import("cupy")
+from tests.utils import HAS_CUPY, assert_allclose, skip_if_no_cuda
 
 
 class TestToNumpyd(unittest.TestCase):
-    @skipUnless(has_cp, "CuPy is required.")
+    @skipUnless(HAS_CUPY, "CuPy is required.")
     def test_cupy_input(self):
         test_data = cp.array([[1, 2], [3, 4]])
         test_data = cp.rot90(test_data)
diff --git a/tests/test_to_tensor.py b/tests/test_to_tensor.py
@@ -15,9 +15,7 @@
 from parameterized import parameterized
 
 from monai.transforms import ToTensor
-from tests.utils import TEST_NDARRAYS, assert_allclose, optional_import
-
-cp, has_cp = optional_import("cupy")
+from tests.utils import HAS_CUPY, TEST_NDARRAYS, assert_allclose
 
 im = [[1, 2], [3, 4]]
 
@@ -47,7 +45,7 @@ def test_single_input(self, test_data):
         assert_allclose(result, test_data, type_test=False)
         self.assertEqual(result.ndim, 0)
 
-    @unittest.skipUnless(has_cp, "CuPy is required.")
+    @unittest.skipUnless(HAS_CUPY, "CuPy is required.")
     def test_cupy(self):
         test_data = [[1, 2], [3, 4]]
         cupy_array = cp.ascontiguousarray(cp.asarray(test_data))
diff --git a/tests/utils.py b/tests/utils.py
@@ -214,6 +214,27 @@ def __call__(self, obj):
         )(obj)
 
 
+def has_cupy():
+    """
+    Returns True if the user has installed a version of cupy.
+    """
+    cp, has_cp = optional_import("cupy")
+    if not has_cp:
+        return False
+    try:
+        x = cp.arange(6, dtype="f").reshape(2, 3)
+        y = cp.arange(3, dtype="f")
+        kernel = cp.ElementwiseKernel(
+            "float32 x, float32 y", "float32 z", """ if (x - 2 > y) { z = x * y; } else { z = x + y; } """, "my_kernel"
+        )
+        return kernel(x, y)[0, 0] == 0
+    except Exception:
+        return False
+
+
+HAS_CUPY = has_cupy()
+
+
 def make_nifti_image(array: NdarrayOrTensor, affine=None):
     """
     Create a temporary nifti image on the disk and return the image name.
