Iterate on multi GPU

python/clij2fft/richardson_lucy_dask_multi_gpu.py

@@ -141,30 +141,54 @@ def richardson_lucy_dask(img, psf, numiterations, regularizationfactor, non_circ
 
     chunk_size = (img.shape[0], y_chunk_size, x_chunk_size)
     print('chunk size is',chunk_size)
+    print('==========================================================================')
 
     dimg = da.from_array(img,chunks=(img.shape[0], y_chunk_size, x_chunk_size))
 
+    from multiprocessing import Pool, current_process, Queue
+    queue = Queue()
+
+    for i in range(1):
+        queue.put(i)
+
     #from dask.distributed import get_worker
     if non_circulant:
 
         def richardson_lucy_nc_dask_task(img, psf, numiterations, regularizationfactor=0, lib=None, block_info=None, block_id=None, thread_id=None):
-            print('block id is', block_id)
-            print('block info is', block_info)
-            print('thread id is', thread_id)
-            #print('worker is ', get_worker())
-            return richardson_lucy_nc(img, psf, numiterations, regularizationfactor=regularizationfactor, lib=lib)
+
+            try:
+                print()
+                gpu_num=queue.get()
+                print('start rlnc')
+                print('gpu num is', gpu_num)
+                print('block id is', block_id)
+                print('block info is', block_info)
+                print('thread id is', thread_id)
+                #print('worker is ', get_worker())
+                if block_id is None:
+                    print('returning block id is None')
+                    return None
+                result=richardson_lucy_nc(img, psf, numiterations, regularizationfactor=regularizationfactor, lib=lib)
+                print('end rlnc')
+                return result
+            except:
+                pass
+            finally:
+                print('putting gpu num back', gpu_num)
+                queue.put(gpu_num)
         rl_func = richardson_lucy_nc_dask_task
     else:
         def richardson_lucy_dask_task(img, psf, numiterations, regularizationfactor=0, lib=None, block_info=None, block_id=None, thread_id=None):
-            print('block id is', block_id)
+            print('\nblock id is', block_id)
             print('block info is', block_info)
             print('thread id is', thread_id)
             #print('worker is', get_worker())
-            return richardson_lucy(img, psf, numiterations, regularizationfactor=regularizationfactor, lib=lib)
+            return img#richardson_lucy(img, psf, numiterations, regularizationfactor=regularizationfactor, lib=lib)
         rl_func = richardson_lucy_dask_task
 
+
     out = dimg.map_overlap(rl_func, depth={0:0, 1:overlap, 2:overlap}, dtype=np.float32, psf=psf, numiterations=numiterations, regularizationfactor=regularizationfactor)
-    return out.compute(num_workers=1)
+    return out.compute(num_workers=4)
 
 
 

python/clij2fft/test_richardson_lucy_dask.py

@@ -1,24 +1,33 @@
-from richardson_lucy_dask import richardson_lucy_dask
+from richardson_lucy_dask_multi_gpu import richardson_lucy_dask
 from skimage.io import imread
 import numpy as np
 from matplotlib import pyplot as plt
 
 img_name=r'D:\\images/images/Bars-G10-P15-stack-cropped.tif'
 psf_name=r'D:\\images/images/PSF-Bars-stack-cropped.tif'
 
+img_name=r'/home/bnorthan/images/deconvolution/Bars-G10-P15-stack.tif'
+psf_name=r'/home/bnorthan/images/deconvolution/PSF-Bars-stack.tif'
+
 img=imread(img_name)
 print('image shape is',img.shape)
 
 pad_z=50
-pad_y=291
-pad_x=700
-mem_to_use=8
+pad_y=50
+pad_x=50
+mem_to_use=1
 
 img = np.pad(img, [(pad_z,pad_z),(pad_y, pad_y),(pad_x, pad_x)], mode = 'constant', constant_values = 0)
 print('image shape is',img.shape)
 psf=imread(psf_name)
 
-decon=richardson_lucy_dask(img, psf, 100, 0.0001, mem_to_use=mem_to_use)
+decon=richardson_lucy_dask(img, psf, 10, 0.0001, mem_to_use=mem_to_use)
+
+fig, ax = plt.subplots(1,2)
+ax[0].imshow(img.max(axis=0))
+ax[0].set_title('img')
+
+ax[1].imshow(decon.max(axis=0))
+ax[1].set_title('deconvolution')
 
-plt.imshow(decon.max(axis=0))
-plt.show()
+plt.show()