Parallelization problem for 3D tensor #598

Activity

1. rflamary

   added

   question

   on Jan 15, 2024
2. rflamary commented on Jan 15, 2024

   

   rflamary

   on Jan 15, 2024

   Collaborator

   More actions

   Hello EMD is already optimized C++ and is not paralelizable directly but you can use jobli to comptue multiple emd in parallel. here is some code that was sent to me by the great @tomMoral

   # %%
   
   import numpy as np
   import ot
   
   from joblib import Parallel, delayed
   
   # %%
   
   n = 1000
   k = 50
   
   M = np.random.rand(k, n, n)
   a = np.ones(n) / n
   
   # %% loop
   
   
   def emd(M, axes=None):
       return ot.emd(a, a, M)
   
   
   R_loop = np.zeros((k, n, n))
   ot.tic()
   for i in range(k):
       R_loop[i] = emd(M[i])
   ot.toc()
   
   
   # %% numpy take/stack
   
   def apply_across_axis(func, M, axis=0):
       return np.stack([
           func(M.take(i, axis))
           for i in range(M.shape[axis])
       ], axis=axis)
   
   
   ot.tic()
   R_numpy = apply_across_axis(emd, M, 0)
   ot.toc()
   
   
   # %% joblib?
   
   def apply_across_axis_joblib(func, M, axis=0, n_jobs=4):
       res = Parallel(n_jobs=n_jobs, max_nbytes=None)(
           delayed(func)(M.take(i, axis))
           for i in range(M.shape[axis])
       )
       return np.stack(res, axis=axis)
   
   
   R_joblib = apply_across_axis_joblib(emd, M[:4], 0)
   ot.tic()
   R_joblib = apply_across_axis_joblib(emd, M, 0)
   ot.toc()

   I think it can be easily adapted to your problem (I dont know what is going to happen with torch tensors though)
3. PythonOT

   locked and limited conversation to collaborators on Jan 15, 2024
4. rflamary

   converted this issue into a discussion #600 on Jan 15, 2024