{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "\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" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "# Author: Paul Krzakala \n# License: MIT License\n\n# sphinx_gallery_thumbnail_number = 1" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 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

Note

A straightforward approach would be to use a Python loop and\n :func:`ot.dist`.\n However, this is inefficient when working with batches.

\n\nInstead, you can directly use :func:`ot.batch.dist_batch`, which computes\nall cost matrices in parallel.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "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.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])" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Solving the Problems\n\nOnce the cost matrices are computed, we can solve the corresponding\noptimal transport problems.\n\n

Note

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.

\n\nInstead, you can use :func:`ot.batch.solve_batch`, which solves all\nproblems in parallel.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "reg = 1.0\nmax_iter = 100\ntol = 1e-3\n\n# Naive approach\nresults_values_list = []\nfor i in range(n_problems):\n res = ot.solve(M_list[i], reg=reg, max_iter=max_iter, tol=tol, reg_type=\"entropy\")\n results_values_list.append(res.value_linear)\n\n# Batched approach\nresults_batch = ot.solve_batch(\n M=M_batch, reg=reg, max_iter=max_iter, tol=tol, reg_type=\"entropy\"\n)\nresults_values_batch = results_batch.value_linear\n\nassert np.allclose(np.array(results_values_list), results_values_batch, atol=tol * 10)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Comparing Computation Time\n\nWe now compare the runtime of the two approaches on larger problems.\n\n

Note

The speedup obtained with :mod:`ot.batch` can be even more\n significant when computations are performed on a GPU.

\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "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.dist_batch(samples_source, samples_target)\n res_batch = ot.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\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Gromov-Wasserstein\n\nThe :mod:`ot.batch` module also provides a batched Gromov-Wasserstein solver.\n\n

Note

This solver is **not** equivalent to calling :func:`ot.solve_gromov`\n repeatedly in a loop.

\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" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "from ot import solve_gromov, solve_gromov_batch\n\n\ndef benchmark_naive_gw(samples_source, samples_target):\n start = perf_counter()\n avg_value = 0\n for i in range(n_problems):\n C1 = ot.dist(samples_source[i], samples_source[i])\n C2 = ot.dist(samples_target[i], samples_target[i])\n res = solve_gromov(C1, C2, max_iter=1000, tol=tol)\n avg_value += res.value\n avg_value /= n_problems\n end = perf_counter()\n return end - start, avg_value\n\n\ndef benchmark_batch_gw(samples_source, samples_target):\n start = perf_counter()\n C1_batch = ot.dist_batch(samples_source, samples_source)\n C2_batch = ot.dist_batch(samples_target, samples_target)\n res_batch = solve_gromov_batch(\n C1_batch, C2_batch, reg=1, max_iter=100, max_iter_inner=50, tol=tol\n )\n avg_value = np.mean(res_batch.value)\n end = perf_counter()\n return end - start, avg_value\n\n\ntime_naive_gw, avg_value_naive_gw = benchmark_naive_gw(samples_source, samples_target)\ntime_batch_gw, avg_value_batch_gw = benchmark_batch_gw(samples_source, samples_target)\n\nprint(f\"{'Method':<20}{'Time (s)':<15}{'Avg Value':<15}\")\nprint(f\"{'Naive GW':<20}{time_naive_gw:<15.4f}{avg_value_naive_gw:<15.4f}\")\nprint(f\"{'Batched GW':<20}{time_batch_gw:<15.4f}{avg_value_batch_gw:<15.4f}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## In summary: no more for loops!\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import matplotlib.pyplot as plt\n\nfig, ax = plt.subplots(figsize=(4, 4))\nax.text(0.5, 0.5, \"For\", fontsize=160, ha=\"center\", va=\"center\", zorder=0)\nax.axis(\"off\")\nax.plot([0, 1], [0, 1], color=\"red\", linewidth=10, zorder=1)\nax.plot([0, 1], [1, 0], color=\"red\", linewidth=10, zorder=1)\nplt.show()" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.10.18" } }, "nbformat": 4, "nbformat_minor": 0 }