sinkhorn GPU implementation

1 files changed, 81 insertions, 0 deletions

diff --git a/ot/gpu/da.py b/ot/gpu/da.py
new file mode 100644
index 0000000..7d04b5b
--- /dev/null
+++ b/ot/gpu/da.py
@@ -0,0 +1,81 @@
+# -*- coding: utf-8 -*-
+"""
+Domain adaptation with optimal transport and GPU
+"""
+
+import numpy as np
+from ..utils import unif
+from ..da import OTDA
+from .bregman import sinkhornGPU
+
+
+def pairwiseEuclideanGPU(a, b, returnAsGPU=False, squared=False, cudamat=None):
+    # a is shape (n, f) and b shape (m, f). Return matrix c of shape (n, m).
+    # First compute in c_GPU the squared euclidean distance. And return its
+    # square root. At each cell [i,j] of c, we want to have
+    # sum{k in range(f)} ( (a[i,k] - b[j,k])^2 ). We know that
+    # (a-b)^2 = a^2 -2ab +b^2. Thus we want to have in each cell of c:
+    # sum{k in range(f)} ( a[i,k]^2 -2a[i,k]b[j,k] +b[j,k]^2).
+
+    a_GPU = cudamat.CUDAMatrix(a)
+    b_GPU = cudamat.CUDAMatrix(b)
+
+    # Multiply a by b transpose to obtain in each cell [i,j] of c the
+    # value sum{k in range(f)} ( a[i,k]b[j,k] )
+    c_GPU = cudamat.dot(a_GPU, b_GPU.transpose())
+    # multiply by -2 to have sum{k in range(f)} ( -2a[i,k]b[j,k] )
+    c_GPU.mult(-2)
+
+    # Compute the vectors of the sum of squared elements.
+    a_GPU = cudamat.pow(a_GPU, 2).sum(axis=1)
+    b_GPU = cudamat.pow(b_GPU, 2).sum(axis=1)
+
+    # Add the vectors in each columns (respectivly rows) of c.
+    # sum{k in range(f)} ( a[i,k]^2 -2a[i,k]b[j,k] )
+    c_GPU.add_col_vec(a_GPU)
+    # sum{k in range(f)} ( a[i,k]^2 -2a[i,k]b[j,k] +b[j,k]^2)
+    c_GPU.add_row_vec(b_GPU.transpose())
+
+    if not squared:
+        c_GPU = cudamat.sqrt(c_GPU)
+
+    if returnAsGPU:
+        return c_GPU
+    else:
+        return c_GPU.asarray()
+
+
+class OTDA_sinkhorn_GPU(OTDA):
+    def fit(self, xs, xt, reg=1, ws=None, wt=None, norm=None):
+        import cudamat
+        cudamat.init()
+        xs = np.asarray(xs, dtype=np.float64)
+        xt = np.asarray(xt, dtype=np.float64)
+
+        self.xs = xs
+        self.xt = xt
+
+        if wt is None:
+            wt = unif(xt.shape[0])
+        if ws is None:
+            ws = unif(xs.shape[0])
+
+        self.ws = ws
+        self.wt = wt
+
+        self.M_GPU = pairwiseEuclideanGPU(xs, xt, returnAsGPU=True,
+                                          squared=True, cudamat=cudamat)
+
+        if norm == "median":
+            self.M_GPU.divide(float(np.median(self.M_GPU.asarray())))
+        elif norm == "max":
+            self.M_GPU.divide(float(np.max(self.M_GPU.asarray())))
+        elif norm == "log":
+            M = np.log(1 + self.M_GPU.asarray())
+            self.M_GPU = cudamat.CUDAMatrix(M)
+        elif norm == "loglog":
+            M = np.log(1 + np.log(1 + self.M_GPU.asarray()))
+            self.M_GPU = cudamat.CUDAMatrix(M)
+
+        self.G = sinkhornGPU(ws, wt, self.M_GPU, reg, cudamat=cudamat)
+        self.computed = True