Batch OT losses (Sinkhorn + Gromov)

RELEASES.md

@@ -19,6 +19,7 @@
 - Fix jax version for auto-grad (PR #732)
 - Added to each example in the examples gallery the information about the release version in which it was introduced (PR #743)
 - Removed release information from quickstart guide (PR #744)
+- Implement batch parallel solvers in ot.batch (PR #745)
 
 #### Closed issues
 - Fixed `ot.mapping` solvers which depended on deprecated `cvxpy` `ECOS` solver (PR #692, Issue #668)

examples/batch/README.txt

@@ -0,0 +1,4 @@
+
+
+Batch parrallel optimal transport
+-----------------------------------

examples/batch/demo_efficiency.py

@@ -0,0 +1,114 @@
+# -*- coding: utf-8 -*-
+"""
+==================
+Batch parallel OT
+==================
+
+Shows the efficiency of using parallel OT solvers for optimal transport.
+
+"""
+
+# Author: Paul Krzakala <paul.krzakala@gmail.com>
+# License: MIT License
+
+from ot.batch._linear import *
+from ot.batch._quadratic import *
+import numpy as np
+import matplotlib.pyplot as plt
+from time import perf_counter
+import torch
+
+b = 256
+n = 40
+d_nodes = 16
+d_edges = 16
+epsilon = 0.1
+max_iter = 10
+tol = 1e-5
+max_iter_inner = 10
+tol_inner = 1e-5
+alpha = 0.5
+
+X1 = torch.randn(b, n, d_nodes, dtype=torch.float32)
+C1 = torch.randn(b, n, n, d_edges, dtype=torch.float32)
+
+X2 = X1[:, torch.randperm(n), :]
+perm = torch.randperm(n)
+C2 = C1[:, perm, :, :][:, :, perm, :]
+
+M = cost_matrix_l2_batch(X1, X2)
+
+# Compute using for loops and batches of size 1
+start = perf_counter()
+average_cost = 0
+for i in range(b):
+    res = solve_gromov_batch(
+        alpha=alpha,
+        epsilon=epsilon,
+        M=M[i : i + 1],
+        C1=C1[i : i + 1],
+        C2=C2[i : i + 1],
+        max_iter=max_iter,
+        tol=tol,
+        max_iter_inner=max_iter_inner,
+        tol_inner=tol_inner,
+    )
+    average_cost += res.value.item()
+average_cost /= b
+
+end = perf_counter()
+print(f"Quadratic solver naive (CPU): {end - start:.4f} seconds")
+print(f"Average cost (CPU): {average_cost:.4f}")
+print("")
+
+start = perf_counter()
+res = solve_gromov_batch(
+    alpha=alpha,
+    epsilon=epsilon,
+    M=M,
+    C1=C1,
+    C2=C2,
+    max_iter=max_iter,
+    tol=tol,
+    max_iter_inner=max_iter_inner,
+    tol_inner=tol_inner,
+)
+average_cost = res.value.mean().item()
+end = perf_counter()
+print(f"Quadratic solver batch (CPU): {end - start:.4f} seconds")
+print(f"Average cost (CPU): {average_cost:.4f}")
+print("")
+
+if torch.cuda.is_available():
+    M, C1, C2 = M.cuda(), C1.cuda(), C2.cuda()
+    # GPU Warmup
+    print("Warming up GPU...")
+    for _ in range(100):
+        solve_gromov_batch(
+            alpha=alpha,
+            epsilon=epsilon,
+            M=M,
+            C1=C1,
+            C2=C2,
+            max_iter=max_iter,
+            tol=tol,
+            max_iter_inner=max_iter_inner,
+            tol_inner=tol_inner,
+        )
+    print("Done.")
+    start = perf_counter()
+    out = solve_gromov_batch(
+        alpha=alpha,
+        epsilon=epsilon,
+        M=M,
+        C1=C1,
+        C2=C2,
+        max_iter=max_iter,
+        tol=tol,
+        max_iter_inner=max_iter_inner,
+        tol_inner=tol_inner,
+    )
+    average_cost = out["cost"].mean().item()
+    end = perf_counter()
+    print(f"Quadratic solver batch (GPU): {end - start:.4f} seconds")
+    print(f"Average cost (GPU): {average_cost:.4f}")

examples/batch/demo_pitfalls.py

@@ -0,0 +1,136 @@
+# -*- coding: utf-8 -*-
+"""
+==================
+Batch parallel linear OT
+==================
+
+Showcase the importance of correctly setting the "symmetric" boolean.
+
+"""
+
+from ot.batch._linear import *
+from ot.batch._quadratic import *
+import numpy as np
+import torch
+from time import perf_counter
+
+np.random.seed(0)
+
+
+def demo_symmetric():
+    print("")
+    print("Generate random data with C1 and C2 NOT symmetric...")
+    b = 256
+    n = 40
+    epsilon = 0.1
+    max_iter = 10
+    tol = 1e-5
+    max_iter_inner = 10
+    tol_inner = 1e-5
+
+    C1 = torch.randn(b, n, n, dtype=torch.float32)
+
+    perm = torch.randperm(n)
+    C2 = C1[:, perm, :][:, :, perm]
+    C2 = torch.randn(b, n, n, dtype=torch.float32)
+
+    print("")
+
+    print("Solving gromov with symmetric=True...")
+    start = perf_counter()
+    res = solve_gromov_batch(
+        epsilon=epsilon,
+        C1=C1,
+        C2=C2,
+        max_iter=max_iter,
+        tol=tol,
+        max_iter_inner=max_iter_inner,
+        tol_inner=tol_inner,
+        symmetric=True,
+    )
+    end = perf_counter()
+    n_iter = res.log["n_iter"]
+    print(f"... solver took: {n_iter} iterations.")
+    print(f"... solver took: {end - start:.4f} seconds")
+    print(f"... final value: {res.value.mean().item():.4f}")
+
+    print("")
+
+    print("Solving gromov with symmetric=False...")
+    start = perf_counter()
+    res = solve_gromov_batch(
+        epsilon=epsilon,
+        C1=C1,
+        C2=C2,
+        max_iter=max_iter,
+        tol=tol,
+        max_iter_inner=max_iter_inner,
+        tol_inner=tol_inner,
+        symmetric=False,
+    )
+    end = perf_counter()
+    n_iter = res.log["n_iter"]
+    print(f"... solver took: {n_iter} iterations.")
+    print(f"... solver took: {end - start:.4f} seconds")
+    print(f"... final value: {res.value.mean().item():.4f}")
+
+    print("")
+    print(
+        "Conclusion: symmetric=False is slightly faster but leads to suboptimal values because some incorrect approximations are made."
+    )
+
+
+def demo_recompute_const():
+    print("")
+    print("Generate random data...")
+    b = 64
+    n = 64
+    d_nodes = 4
+    d_edges = 2
+
+    C1 = np.random.randn(b, n, n, d_edges).astype("float32")
+    C2 = np.random.randn(b, n, n, d_edges).astype("float32")
+
+    p = np.ones((b, n), dtype=np.float32) / n
+    q = np.ones((b, n), dtype=np.float32) / n
+
+    def compute_loss(T, recompute_const=False):
+        metric = QuadraticEuclidean()
+        L = metric.cost_tensor(p, q, C1, C2, symmetric=True)
+        LT = tensor_product(L, T, recompute_const=recompute_const, symmetric=True)
+        loss = (LT * T).sum((1, 2))
+        return loss
+
+    print("")
+    print("Generate random coupling matrix T that does not satisfy marginals...")
+    T = np.random.rand(b, n, n).astype("float32")
+
+    print("")
+    print("Compute the gromov loss - recompute_const=False")
+    loss = compute_loss(T, recompute_const=False)
+    is_pos = (loss >= -1e-6).all()
+    print("All losses are positive: ", is_pos)
+
+    print("")
+    print("Compute the gromov loss - recompute_const=True")
+    T = np.random.rand(b, n, n).astype("float32")
+    loss = compute_loss(T, recompute_const=True)
+    is_pos = (loss >= -1e-6).all()
+    print("All losses are positive: ", is_pos)
+
+    print("")
+    print(
+        "Conclusion: when T does not satisfy marginals, it is important to set recompute_const=True to ensure that the loss computation is correct."
+    )
+
+
+if __name__ == "__main__":
+    print(
+        "----------------- Setting symmetric Flag correctly in ot.solve_gromov  -----------------"
+    )
+    demo_symmetric()
+    print("")
+    print(
+        "----------------- Testing recompute_const option in tensor_product  -----------------"
+    )
+    demo_recompute_const()

examples/batch/demo_solve_batch.py

@@ -0,0 +1,61 @@
+# -*- coding: utf-8 -*-
+"""
+==================
+Batch parallel linear OT
+==================
+
+Shows how to use the batch parallel linear OT solvers.
+
+"""
+
+# Author: Paul Krzakala <paul.krzakala@gmail.com>
+# License: MIT License
+
+from ot.batch._linear import *
+from ot.batch._quadratic import *
+import numpy as np
+import matplotlib.pyplot as plt
+
+# Set parameters
+b = 4
+n = 20
+d = 4
+noise = 0.1
+max_iter = 10000
+tol = 1e-5
+
+# Generate random data
+X = np.random.randn(b, n, d).astype("float32")
+Y = np.random.randn(b, n, d).astype("float32")
+M = cost_matrix_l2_batch(X, Y)
+
+# Define grid search
+epsilons = np.logspace(-2, 1, num=10)
+
+plt.figure(figsize=(10, 6))
+
+cost_list = []
+n_iter_list = []
+for epsilon in epsilons:
+    res = solve_batch(epsilon=epsilon, M=M, max_iter=max_iter, tol=tol, log_dual=True)
+    n_iter = res.log["n_iter"]
+    cost = res.value_linear.mean()
+    cost_list.append(cost)
+    n_iter_list.append(n_iter)
+
+ax1 = plt.gca()
+ax2 = ax1.twinx()
+
+ax1.plot(epsilons, n_iter_list, "r-o", label="Iterations")
+ax2.plot(epsilons, cost_list, "b-s", label="Cost")
+
+ax1.set_xlabel("Epsilon")
+ax1.set_ylabel("Iterations", color="r")
+ax2.set_ylabel("Cost", color="b")
+ax1.set_xscale("log")
+
+ax1.tick_params(axis="y", labelcolor="r")
+ax2.tick_params(axis="y", labelcolor="b")
+
+plt.title("Epsilon vs Iters and Cost")
+plt.show()

examples/batch/demo_solve_gromov_batch.py

@@ -0,0 +1,64 @@
+# -*- coding: utf-8 -*-
+"""
+==================
+Batch parallel linear OT
+==================
+
+Shows how to use the batch parallel quadratic OT solvers.
+
+"""
+
+from ot.batch._linear import *
+from ot.batch._quadratic import *
+import numpy as np
+import matplotlib.pyplot as plt
+
+b = 32
+n = 20
+d_nodes = 4
+d_edges = 2
+epsilon = 0.1
+max_iter = 10
+tol = 1e-5
+max_iter_inner = 10
+tol_inner = 1e-5
+
+X1 = np.random.randn(b, n, d_nodes).astype("float32")
+C1 = np.random.randn(b, n, n, d_edges).astype("float32")
+
+X2 = X1[:, np.random.permutation(n), :]
+permutation = np.random.permutation(n)
+C2 = C1[:, permutation, :, :][:, :, permutation, :]
+
+M = cost_matrix_l2_batch(X1, X2)
+
+cost_linear_list = []
+cost_quadratic_list = []
+alpha_list = np.linspace(0.001, 0.999, 20)
+
+for alpha in alpha_list:
+    res = solve_gromov_batch(
+        alpha=alpha,
+        epsilon=epsilon,
+        M=M,
+        C1=C1,
+        C2=C2,
+        max_iter=max_iter,
+        tol=tol,
+        max_iter_inner=max_iter_inner,
+        tol_inner=tol_inner,
+    )
+
+    cost_linear = res.value_linear.mean() / (1 - alpha)
+    cost_quadratic = res.value_quad.mean() / (alpha)
+    cost_linear_list.append(cost_linear)
+    cost_quadratic_list.append(cost_quadratic)
+
+plt.figure(figsize=(10, 6))
+plt.plot(alpha_list, cost_linear_list, marker="o", label="Cost Linear")
+plt.plot(alpha_list, cost_quadratic_list, marker="o", label="Cost Quadratic")
+plt.xlabel("Alpha")
+plt.ylabel("Average Cost")
+plt.legend()
+plt.title("Cost vs Alpha")
+plt.show()