summaryrefslogtreecommitdiff
path: root/ot/utils.py
diff options
context:
space:
mode:
authorLeoGautheron <leo_g_autheron@hotmail.fr>2018-07-11 22:28:38 +0200
committerLeoGautheron <leo_g_autheron@hotmail.fr>2018-07-11 22:28:38 +0200
commitcb6bdc516697e3bad6776b897f22c8b6a22f13cd (patch)
tree7b040dc809893613390d202b41e539cc69972fe7 /ot/utils.py
parent39cbcd302c1d1e275c628d3bac073ec1f89596c6 (diff)
Speed-up Sinkhorn
Speed-up in 3 places: - the computation of pairwise distance is faster with sklearn.metrics.pairwise.euclidean_distances - faster computation of K = np.exp(-M / reg) - faster computation of the error every 10 iterations Example with this little script: import time import numpy as np import ot rng = np.random.RandomState(0) transport = ot.da.SinkhornTransport() time1 = time.time() Xs, ys, Xt = rng.randn(10000, 100), rng.randint(0, 2, size=10000), rng.randn(10000, 100) transport.fit(Xs=Xs, Xt=Xt) time2 = time.time() print("OT Computation Time {:6.2f} sec".format(time2-time1)) transport = ot.da.SinkhornLpl1Transport() transport.fit(Xs=Xs, ys=ys, Xt=Xt) time3 = time.time() print("OT LpL1 Computation Time {:6.2f} sec".format(time3-time2)) Before OT Computation Time 19.93 sec OT LpL1 Computation Time 133.43 sec After OT Computation Time 7.55 sec OT LpL1 Computation Time 82.25 sec
Diffstat (limited to 'ot/utils.py')
-rw-r--r--ot/utils.py4
1 files changed, 3 insertions, 1 deletions
diff --git a/ot/utils.py b/ot/utils.py
index 7dac283..5b052ac 100644
--- a/ot/utils.py
+++ b/ot/utils.py
@@ -13,6 +13,7 @@ import time
import numpy as np
from scipy.spatial.distance import cdist
+from sklearn.metrics.pairwise import euclidean_distances
import sys
import warnings
try:
@@ -104,7 +105,8 @@ def dist(x1, x2=None, metric='sqeuclidean'):
"""
if x2 is None:
x2 = x1
-
+ if metric == "sqeuclidean":
+ return euclidean_distances(x1, x2, squared=True)
return cdist(x1, x2, metric=metric)