Some speed tests for you @tom, only x10 at the moment?
It should be faster when i) it’s run on the gpu, ii) the histograms get bigger, iii) increase the dimensions, i.e images or higher dimension tensors, iv) tune lambda
Check out figure 4, from Marco Cuturi’s original paper,

  
      papers.nips.cc
  
  
    
4927-sinkhorn-distances-lightspeed-computation-of-optimal-transport.pdf 34

360.42 KB

  
  
    
    
  
  


#-----------------------------------------------------------------------------------------
#%% Speed Tests - see https://github.com/rflamary/POT/blob/master/examples/plot_OT_1D.py
#-----------------------------------------------------------------------------------------

import numpy as np
import pyemd 
import ot
import functools
from time import time
from ot.datasets import get_1D_gauss as gauss

#@functools.lru_cache (maxsize=None) # hash(tuple( yadda, yadda
def pyemd_linprog(n,a,b,M):
    start = time()
    for i in range(n): 
        div_pyemd = pyemd.emd(a,b,M)
    run_time = time() - start
    print(run_time)

#@functools.lru_cache (maxsize=None)
def pot_sinkhorn(n,a,b,M,lambd):
    start = time()
    for i in range(n): 
        T_sink = ot.sinkhorn(a,b,M, lambd)
        #div_pot_sink = np.trace( np.matmul( T_sink.transpose() , M) )
    run_time = time() - start
    print(run_time)

#%% parameters

n=128 # nb bins

# bin positions
x=np.arange(n,dtype=np.float64)

# Gaussian distributions
a=gauss(n,m=20,s=5) # m= mean, s= std
b=gauss(n,m=60,s=10)

# loss matrix
M=ot.dist(x.reshape((n,1)),x.reshape((n,1)))
M/=M.max()

##--------
runs=1000
lambd = 1

>>> pyemd_linprog( runs, a, b, M )
21.02236819267273
>>> pot_sinkhorn( runs, a, b, M, lambd )
2.0207273960113525