You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
"\n# Solving Many Optimal Transport Problems in Parallel\n\nIn some situations, one may want to solve many OT problems with the same\nstructure (same number of samples, same cost function, etc.) at the same time.\n\nIn that case using a for loop to solve the problems sequentially is inefficient.\nThis example shows how to use the batch solvers implemented in POT to solve\nmany problems in parallel on CPU or GPU (even more efficient on GPU).\n"
8
+
]
9
+
},
10
+
{
11
+
"cell_type": "code",
12
+
"execution_count": null,
13
+
"metadata": {
14
+
"collapsed": false
15
+
},
16
+
"outputs": [],
17
+
"source": [
18
+
"# Author: Paul Krzakala <paul.krzakala@gmail.com>\n# License: MIT License\n\n# sphinx_gallery_thumbnail_number = 1"
19
+
]
20
+
},
21
+
{
22
+
"cell_type": "markdown",
23
+
"metadata": {},
24
+
"source": [
25
+
"## Computing the Cost Matrices\n\nWe want to create a batch of optimal transport problems with\n$n$ samples in $d$ dimensions.\n\nTo do this, we first need to compute the cost matrices for each problem.\n\n<div class=\"alert alert-info\"><h4>Note</h4><p>A straightforward approach would be to use a Python loop and\n :func:`ot.dist`.\n However, this is inefficient when working with batches.</p></div>\n\nInstead, you can directly use :func:`ot.batch.dist_batch`, which computes\nall cost matrices in parallel.\n\n"
26
+
]
27
+
},
28
+
{
29
+
"cell_type": "code",
30
+
"execution_count": null,
31
+
"metadata": {
32
+
"collapsed": false
33
+
},
34
+
"outputs": [],
35
+
"source": [
36
+
"import ot\nimport numpy as np\n\nn_problems = 4 # nb problems/batch size\nn_samples = 8 # nb samples\ndim = 2 # nb dimensions\n\nnp.random.seed(0)\nsamples_source = np.random.randn(n_problems, n_samples, dim)\nsamples_target = samples_source + 0.1 * np.random.randn(n_problems, n_samples, dim)\n\n# Naive approach\nM_list = []\nfor i in range(n_problems):\n M_list.append(\n ot.dist(samples_source[i], samples_target[i])\n ) # List of cost matrices n_samples x n_samples\n# Batched approach\nM_batch = ot.batch.dist_batch(\n samples_source, samples_target\n) # Array of cost matrices n_problems x n_samples x n_samples\n\nfor i in range(n_problems):\n assert np.allclose(M_list[i], M_batch[i])"
37
+
]
38
+
},
39
+
{
40
+
"cell_type": "markdown",
41
+
"metadata": {},
42
+
"source": [
43
+
"## Solving the Problems\n\nOnce the cost matrices are computed, we can solve the corresponding\noptimal transport problems.\n\n<div class=\"alert alert-info\"><h4>Note</h4><p>One option is to solve them sequentially with a Python loop using\n :func:`ot.solve`.\n This is simple but inefficient for large batches.</p></div>\n\nInstead, you can use :func:`ot.batch.solve_batch`, which solves all\nproblems in parallel.\n\n"
"## Comparing Computation Time\n\nWe now compare the runtime of the two approaches on larger problems.\n\n<div class=\"alert alert-info\"><h4>Note</h4><p>The speedup obtained with :mod:`ot.batch` can be even more\n significant when computations are performed on a GPU.</p></div>\n\n"
62
+
]
63
+
},
64
+
{
65
+
"cell_type": "code",
66
+
"execution_count": null,
67
+
"metadata": {
68
+
"collapsed": false
69
+
},
70
+
"outputs": [],
71
+
"source": [
72
+
"from time import perf_counter\n\nn_problems = 128\nn_samples = 8\ndim = 2\nreg = 10.0\nmax_iter = 1000\ntol = 1e-3\n\nsamples_source = np.random.randn(n_problems, n_samples, dim)\nsamples_target = samples_source + 0.1 * np.random.randn(n_problems, n_samples, dim)\n\n\ndef benchmark_naive(samples_source, samples_target):\n start = perf_counter()\n for i in range(n_problems):\n M = ot.dist(samples_source[i], samples_target[i])\n res = ot.solve(M, reg=reg, max_iter=max_iter, tol=tol, reg_type=\"entropy\")\n end = perf_counter()\n return end - start\n\n\ndef benchmark_batch(samples_source, samples_target):\n start = perf_counter()\n M_batch = ot.batch.dist_batch(samples_source, samples_target)\n res_batch = ot.batch.solve_batch(\n M=M_batch, reg=reg, max_iter=max_iter, tol=tol, reg_type=\"entropy\"\n )\n end = perf_counter()\n return end - start\n\n\ntime_naive = benchmark_naive(samples_source, samples_target)\ntime_batch = benchmark_batch(samples_source, samples_target)\n\nprint(f\"Naive approach time: {time_naive:.4f} seconds\")\nprint(f\"Batched approach time: {time_batch:.4f} seconds\")"
73
+
]
74
+
},
75
+
{
76
+
"cell_type": "markdown",
77
+
"metadata": {},
78
+
"source": [
79
+
"## Gromov-Wasserstein\n\nThe :mod:`ot.batch` module also provides a batched Gromov-Wasserstein solver.\n\n<div class=\"alert alert-info\"><h4>Note</h4><p>This solver is **not** equivalent to calling :func:`ot.solve_gromov`\n repeatedly in a loop.</p></div>\n\nKey differences:\n\n- :func:`ot.solve_gromov`\n Uses the conditional gradient algorithm. Each inner iteration relies on\n an exact EMD solver.\n\n- :func:`ot.batch.solve_gromov_batch`\n Uses a proximal variant, where each inner iteration applies entropic\n regularization.\n\nAs a result:\n\n- :func:`ot.solve_gromov` is usually faster on CPU\n- :func:`ot.batch.solve_gromov_batch` is slower on CPU, but provides\n better objective values.\n\n.. tip::\n If your data is on a GPU, :func:`ot.batch.solve_gromov_batch`\n is significantly faster AND provides better objective values.\n\n"
0 commit comments