TST: Add explicit test for CUDA recon methods (#492)

tomopy · Jun 12, 2020 · dbf0444 · dbf0444
1 parent 21cca25
commit dbf0444
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 2 deletions.
diff --git a/test/test_tomopy/test_data/mlem_gpu.npy b/test/test_tomopy/test_data/mlem_gpu.npy
diff --git a/test/test_tomopy/test_data/sirt_gpu.npy b/test/test_tomopy/test_data/sirt_gpu.npy
diff --git a/test/test_tomopy/test_recon/test_algorithm.py b/test/test_tomopy/test_recon/test_algorithm.py
@@ -122,9 +122,15 @@ def test_mlem(self):
 
     def test_mlem_accel(self):
         result = recon(self.prj, self.ang, algorithm='mlem', num_iter=4,
-                       accelerated=True)
+                       accelerated=True, device='cpu')
         assert_allclose(result, read_file('mlem_accel.npy'), rtol=1e-2)
 
+    @unittest.skipUnless("CUDA_PATH" in os.environ, "CUDA_PATH not set.")
+    def test_mlem_gpu(self):
+        result = recon(self.prj, self.ang, algorithm='mlem', num_iter=4,
+                       accelerated=True, device='gpu')
+        assert_allclose(result, read_file('mlem_gpu.npy'), rtol=1e-2)
+
     def test_osem(self):
         assert_allclose(
             recon(self.prj, self.ang, algorithm='osem', num_iter=4),
@@ -156,9 +162,15 @@ def test_sirt(self):
 
     def test_sirt_accel(self):
         result = recon(self.prj, self.ang, algorithm='sirt',
-                       num_iter=4, accelerated=True)
+                       num_iter=4, accelerated=True, device='cpu')
         assert_allclose(result, read_file('sirt_accel.npy'), rtol=1e-2)
 
+    @unittest.skipUnless("CUDA_PATH" in os.environ, "CUDA_PATH not set.")
+    def test_sirt_gpu(self):
+        result = recon(self.prj, self.ang, algorithm='sirt',
+                       num_iter=4, accelerated=True, device='gpu')
+        assert_allclose(result, read_file('sirt_gpu.npy'), rtol=1e-2)
+
     def test_tv(self):
         assert_allclose(
             recon(self.prj, self.ang, algorithm='tv', num_iter=4),