diff options
-rw-r--r-- | README.md | 6 | ||||
-rw-r--r-- | docs/source/all.rst | 11 | ||||
-rw-r--r-- | examples/plot_stochastic.py | 207 | ||||
-rw-r--r-- | ot/__init__.py | 2 | ||||
-rw-r--r-- | ot/stochastic.py | 800 | ||||
-rw-r--r-- | test/test_stochastic.py | 191 |
6 files changed, 1216 insertions, 1 deletions
@@ -23,6 +23,7 @@ It provides the following solvers: * Linear OT [14] and Joint OT matrix and mapping estimation [8]. * Wasserstein Discriminant Analysis [11] (requires autograd + pymanopt). * Gromov-Wasserstein distances and barycenters ([13] and regularized [12]) +* Stochastic Optimization for Large-scale Optimal Transport (semi-dual problem [18] and dual problem [19]) Some demonstrations (both in Python and Jupyter Notebook format) are available in the examples folder. @@ -162,6 +163,7 @@ The contributors to this library are: * [Stanislas Chambon](https://slasnista.github.io/) * [Antoine Rolet](https://arolet.github.io/) * Erwan Vautier (Gromov-Wasserstein) +* [Kilian Fatras](https://kilianfatras.github.io/) (Stochastic optimization) This toolbox benefit a lot from open source research and we would like to thank the following persons for providing some code (in various languages): @@ -219,3 +221,7 @@ You can also post bug reports and feature requests in Github issues. Make sure t [16] Agueh, M., & Carlier, G. (2011). [Barycenters in the Wasserstein space](https://hal.archives-ouvertes.fr/hal-00637399/document). SIAM Journal on Mathematical Analysis, 43(2), 904-924. [17] Blondel, M., Seguy, V., & Rolet, A. (2018). [Smooth and Sparse Optimal Transport](https://arxiv.org/abs/1710.06276). Proceedings of the Twenty-First International Conference on Artificial Intelligence and Statistics (AISTATS). + +[18] Genevay, A., Cuturi, M., Peyré, G. & Bach, F. (2016) [Stochastic Optimization for Large-scale Optimal Transport](arXiv preprint arxiv:1605.08527). Advances in Neural Information Processing Systems (2016). + +[19] Seguy, V., Bhushan Damodaran, B., Flamary, R., Courty, N., Rolet, A.& Blondel, M. [Large-scale Optimal Transport and Mapping Estimation](https://arxiv.org/pdf/1711.02283.pdf). International Conference on Learning Representation (2018) diff --git a/docs/source/all.rst b/docs/source/all.rst index bbb9833..9459023 100644 --- a/docs/source/all.rst +++ b/docs/source/all.rst @@ -25,6 +25,11 @@ ot.smooth .. automodule:: ot.smooth :members: +ot.smooth +----- +.. automodule:: ot.smooth + :members: + ot.gromov ---------- @@ -68,3 +73,9 @@ ot.plot .. automodule:: ot.plot :members: + +ot.stochastic +------------- + +.. automodule:: ot.stochastic + :members: diff --git a/examples/plot_stochastic.py b/examples/plot_stochastic.py new file mode 100644 index 0000000..b9375d4 --- /dev/null +++ b/examples/plot_stochastic.py @@ -0,0 +1,207 @@ +""" +========================== +Stochastic examples +========================== + +This example is designed to show how to use the stochatic optimization +algorithms for descrete and semicontinous measures from the POT library. + +""" + +# Author: Kilian Fatras <kilian.fatras@gmail.com> +# +# License: MIT License + +import matplotlib.pylab as pl +import numpy as np +import ot +import ot.plot + + +############################################################################# +# COMPUTE TRANSPORTATION MATRIX FOR SEMI-DUAL PROBLEM +############################################################################# +print("------------SEMI-DUAL PROBLEM------------") +############################################################################# +# DISCRETE CASE +# Sample two discrete measures for the discrete case +# --------------------------------------------- +# +# Define 2 discrete measures a and b, the points where are defined the source +# and the target measures and finally the cost matrix c. + +n_source = 7 +n_target = 4 +reg = 1 +numItermax = 1000 + +a = ot.utils.unif(n_source) +b = ot.utils.unif(n_target) + +rng = np.random.RandomState(0) +X_source = rng.randn(n_source, 2) +Y_target = rng.randn(n_target, 2) +M = ot.dist(X_source, Y_target) + +############################################################################# +# +# Call the "SAG" method to find the transportation matrix in the discrete case +# --------------------------------------------- +# +# Define the method "SAG", call ot.solve_semi_dual_entropic and plot the +# results. + +method = "SAG" +sag_pi = ot.stochastic.solve_semi_dual_entropic(a, b, M, reg, method, + numItermax) +print(sag_pi) + +############################################################################# +# SEMICONTINOUS CASE +# Sample one general measure a, one discrete measures b for the semicontinous +# case +# --------------------------------------------- +# +# Define one general measure a, one discrete measures b, the points where +# are defined the source and the target measures and finally the cost matrix c. + +n_source = 7 +n_target = 4 +reg = 1 +numItermax = 1000 +log = True + +a = ot.utils.unif(n_source) +b = ot.utils.unif(n_target) + +rng = np.random.RandomState(0) +X_source = rng.randn(n_source, 2) +Y_target = rng.randn(n_target, 2) +M = ot.dist(X_source, Y_target) + +############################################################################# +# +# Call the "ASGD" method to find the transportation matrix in the semicontinous +# case +# --------------------------------------------- +# +# Define the method "ASGD", call ot.solve_semi_dual_entropic and plot the +# results. + +method = "ASGD" +asgd_pi, log_asgd = ot.stochastic.solve_semi_dual_entropic(a, b, M, reg, method, + numItermax, log=log) +print(log_asgd['alpha'], log_asgd['beta']) +print(asgd_pi) + +############################################################################# +# +# Compare the results with the Sinkhorn algorithm +# --------------------------------------------- +# +# Call the Sinkhorn algorithm from POT + +sinkhorn_pi = ot.sinkhorn(a, b, M, reg) +print(sinkhorn_pi) + + +############################################################################## +# PLOT TRANSPORTATION MATRIX +############################################################################## + +############################################################################## +# Plot SAG results +# ---------------- + +pl.figure(4, figsize=(5, 5)) +ot.plot.plot1D_mat(a, b, sag_pi, 'semi-dual : OT matrix SAG') +pl.show() + + +############################################################################## +# Plot ASGD results +# ----------------- + +pl.figure(4, figsize=(5, 5)) +ot.plot.plot1D_mat(a, b, asgd_pi, 'semi-dual : OT matrix ASGD') +pl.show() + + +############################################################################## +# Plot Sinkhorn results +# --------------------- + +pl.figure(4, figsize=(5, 5)) +ot.plot.plot1D_mat(a, b, sinkhorn_pi, 'OT matrix Sinkhorn') +pl.show() + + +############################################################################# +# COMPUTE TRANSPORTATION MATRIX FOR DUAL PROBLEM +############################################################################# +print("------------DUAL PROBLEM------------") +############################################################################# +# SEMICONTINOUS CASE +# Sample one general measure a, one discrete measures b for the semicontinous +# case +# --------------------------------------------- +# +# Define one general measure a, one discrete measures b, the points where +# are defined the source and the target measures and finally the cost matrix c. + +n_source = 7 +n_target = 4 +reg = 1 +numItermax = 100000 +lr = 0.1 +batch_size = 3 +log = True + +a = ot.utils.unif(n_source) +b = ot.utils.unif(n_target) + +rng = np.random.RandomState(0) +X_source = rng.randn(n_source, 2) +Y_target = rng.randn(n_target, 2) +M = ot.dist(X_source, Y_target) + +############################################################################# +# +# Call the "SGD" dual method to find the transportation matrix in the +# semicontinous case +# --------------------------------------------- +# +# Call ot.solve_dual_entropic and plot the results. + +sgd_dual_pi, log_sgd = ot.stochastic.solve_dual_entropic(a, b, M, reg, + batch_size, numItermax, + lr, log=log) +print(log_sgd['alpha'], log_sgd['beta']) +print(sgd_dual_pi) + +############################################################################# +# +# Compare the results with the Sinkhorn algorithm +# --------------------------------------------- +# +# Call the Sinkhorn algorithm from POT + +sinkhorn_pi = ot.sinkhorn(a, b, M, reg) +print(sinkhorn_pi) + +############################################################################## +# Plot SGD results +# ----------------- + +pl.figure(4, figsize=(5, 5)) +ot.plot.plot1D_mat(a, b, sgd_dual_pi, 'dual : OT matrix SGD') +pl.show() + + +############################################################################## +# Plot Sinkhorn results +# --------------------- + +pl.figure(4, figsize=(5, 5)) +ot.plot.plot1D_mat(a, b, sinkhorn_pi, 'OT matrix Sinkhorn') +pl.show() diff --git a/ot/__init__.py b/ot/__init__.py index 995ce33..1dde390 100644 --- a/ot/__init__.py +++ b/ot/__init__.py @@ -19,7 +19,7 @@ from . import datasets from . import da from . import gromov from . import smooth - +from . import stochastic # OT functions from .lp import emd, emd2 diff --git a/ot/stochastic.py b/ot/stochastic.py new file mode 100644 index 0000000..5e8206e --- /dev/null +++ b/ot/stochastic.py @@ -0,0 +1,800 @@ +# Author: Kilian Fatras <kilian.fatras@gmail.com> +# +# License: MIT License + +import numpy as np + + +############################################################################## +# Optimization toolbox for SEMI - DUAL problems +############################################################################## + + +def coordinate_grad_semi_dual(b, M, reg, beta, i): + ''' + Compute the coordinate gradient update for regularized discrete + distributions for (i, :) + + The function computes the gradient of the semi dual problem: + + .. math:: + \W_\varepsilon(a, b) = \max_\v \sum_i (\sum_j v_j * b_j + - \reg log(\sum_j exp((v_j - M_{i,j})/reg) * b_j)) * a_i + + where : + - M is the (ns,nt) metric cost matrix + - v is a dual variable in R^J + - reg is the regularization term + - a and b are source and target weights (sum to 1) + + The algorithm used for solving the problem is the ASGD & SAG algorithms + as proposed in [18]_ [alg.1 & alg.2] + + + Parameters + ---------- + + b : np.ndarray(nt,), + target measure + M : np.ndarray(ns, nt), + cost matrix + reg : float nu, + Regularization term > 0 + v : np.ndarray(nt,), + optimization vector + i : number int, + picked number i + + Returns + ------- + + coordinate gradient : np.ndarray(nt,) + + Examples + -------- + + >>> n_source = 7 + >>> n_target = 4 + >>> reg = 1 + >>> numItermax = 300000 + >>> a = ot.utils.unif(n_source) + >>> b = ot.utils.unif(n_target) + >>> rng = np.random.RandomState(0) + >>> X_source = rng.randn(n_source, 2) + >>> Y_target = rng.randn(n_target, 2) + >>> M = ot.dist(X_source, Y_target) + >>> method = "ASGD" + >>> asgd_pi = stochastic.solve_semi_dual_entropic(a, b, M, reg, + method, numItermax) + >>> print(asgd_pi) + + References + ---------- + + [Genevay et al., 2016] : + Stochastic Optimization for Large-scale Optimal Transport, + Advances in Neural Information Processing Systems (2016), + arXiv preprint arxiv:1605.08527. + + ''' + + r = M[i, :] - beta + exp_beta = np.exp(-r / reg) * b + khi = exp_beta / (np.sum(exp_beta)) + return b - khi + + +def sag_entropic_transport(a, b, M, reg, numItermax=10000, lr=None): + ''' + Compute the SAG algorithm to solve the regularized discrete measures + optimal transport max problem + + The function solves the following optimization problem: + + .. math:: + \gamma = arg\min_\gamma <\gamma,M>_F + reg\cdot\Omega(\gamma) + s.t. \gamma 1 = a + \gamma^T 1= b + \gamma \geq 0 + where : + - M is the (ns,nt) metric cost matrix + - :math:`\Omega` is the entropic regularization term + :math:`\Omega(\gamma)=\sum_{i,j} \gamma_{i,j}\log(\gamma_{i,j})` + - a and b are source and target weights (sum to 1) + The algorithm used for solving the problem is the SAG algorithm + as proposed in [18]_ [alg.1] + + + Parameters + ---------- + + a : np.ndarray(ns,), + source measure + b : np.ndarray(nt,), + target measure + M : np.ndarray(ns, nt), + cost matrix + reg : float number, + Regularization term > 0 + numItermax : int number + number of iteration + lr : float number + learning rate + + Returns + ------- + + v : np.ndarray(nt,) + dual variable + + Examples + -------- + + >>> n_source = 7 + >>> n_target = 4 + >>> reg = 1 + >>> numItermax = 300000 + >>> a = ot.utils.unif(n_source) + >>> b = ot.utils.unif(n_target) + >>> rng = np.random.RandomState(0) + >>> X_source = rng.randn(n_source, 2) + >>> Y_target = rng.randn(n_target, 2) + >>> M = ot.dist(X_source, Y_target) + >>> method = "ASGD" + >>> asgd_pi = stochastic.solve_semi_dual_entropic(a, b, M, reg, + method, numItermax) + >>> print(asgd_pi) + + References + ---------- + + [Genevay et al., 2016] : + Stochastic Optimization for Large-scale Optimal Transport, + Advances in Neural Information Processing Systems (2016), + arXiv preprint arxiv:1605.08527. + ''' + + if lr is None: + lr = 1. / max(a / reg) + n_source = np.shape(M)[0] + n_target = np.shape(M)[1] + cur_beta = np.zeros(n_target) + stored_gradient = np.zeros((n_source, n_target)) + sum_stored_gradient = np.zeros(n_target) + for _ in range(numItermax): + i = np.random.randint(n_source) + cur_coord_grad = a[i] * coordinate_grad_semi_dual(b, M, reg, + cur_beta, i) + sum_stored_gradient += (cur_coord_grad - stored_gradient[i]) + stored_gradient[i] = cur_coord_grad + cur_beta += lr * (1. / n_source) * sum_stored_gradient + return cur_beta + + +def averaged_sgd_entropic_transport(a, b, M, reg, numItermax=300000, lr=None): + ''' + Compute the ASGD algorithm to solve the regularized semi contibous measures + optimal transport max problem + + The function solves the following optimization problem: + + .. math:: + \gamma = arg\min_\gamma <\gamma,M>_F + reg\cdot\Omega(\gamma) + s.t. \gamma 1 = a + \gamma^T 1= b + \gamma \geq 0 + where : + - M is the (ns,nt) metric cost matrix + - :math:`\Omega` is the entropic regularization term + :math:`\Omega(\gamma)=\sum_{i,j} \gamma_{i,j}\log(\gamma_{i,j})` + - a and b are source and target weights (sum to 1) + The algorithm used for solving the problem is the ASGD algorithm + as proposed in [18]_ [alg.2] + + + Parameters + ---------- + + b : np.ndarray(nt,), + target measure + M : np.ndarray(ns, nt), + cost matrix + reg : float number, + Regularization term > 0 + numItermax : int number + number of iteration + lr : float number + learning rate + + + Returns + ------- + + ave_v : np.ndarray(nt,) + optimization vector + + Examples + -------- + + >>> n_source = 7 + >>> n_target = 4 + >>> reg = 1 + >>> numItermax = 300000 + >>> a = ot.utils.unif(n_source) + >>> b = ot.utils.unif(n_target) + >>> rng = np.random.RandomState(0) + >>> X_source = rng.randn(n_source, 2) + >>> Y_target = rng.randn(n_target, 2) + >>> M = ot.dist(X_source, Y_target) + >>> method = "ASGD" + >>> asgd_pi = stochastic.solve_semi_dual_entropic(a, b, M, reg, + method, numItermax) + >>> print(asgd_pi) + + References + ---------- + + [Genevay et al., 2016] : + Stochastic Optimization for Large-scale Optimal Transport, + Advances in Neural Information Processing Systems (2016), + arXiv preprint arxiv:1605.08527. + ''' + + if lr is None: + lr = 1. / max(a / reg) + n_source = np.shape(M)[0] + n_target = np.shape(M)[1] + cur_beta = np.zeros(n_target) + ave_beta = np.zeros(n_target) + for cur_iter in range(numItermax): + k = cur_iter + 1 + i = np.random.randint(n_source) + cur_coord_grad = coordinate_grad_semi_dual(b, M, reg, cur_beta, i) + cur_beta += (lr / np.sqrt(k)) * cur_coord_grad + ave_beta = (1. / k) * cur_beta + (1 - 1. / k) * ave_beta + return ave_beta + + +def c_transform_entropic(b, M, reg, beta): + ''' + The goal is to recover u from the c-transform. + + The function computes the c_transform of a dual variable from the other + dual variable: + + .. math:: + u = v^{c,reg} = -reg \sum_j exp((v - M)/reg) b_j + + where : + - M is the (ns,nt) metric cost matrix + - u, v are dual variables in R^IxR^J + - reg is the regularization term + + It is used to recover an optimal u from optimal v solving the semi dual + problem, see Proposition 2.1 of [18]_ + + + Parameters + ---------- + + b : np.ndarray(nt,) + target measure + M : np.ndarray(ns, nt) + cost matrix + reg : float + regularization term > 0 + v : np.ndarray(nt,) + dual variable + + Returns + ------- + + u : np.ndarray(ns,) + + Examples + -------- + + >>> n_source = 7 + >>> n_target = 4 + >>> reg = 1 + >>> numItermax = 300000 + >>> a = ot.utils.unif(n_source) + >>> b = ot.utils.unif(n_target) + >>> rng = np.random.RandomState(0) + >>> X_source = rng.randn(n_source, 2) + >>> Y_target = rng.randn(n_target, 2) + >>> M = ot.dist(X_source, Y_target) + >>> method = "ASGD" + >>> asgd_pi = stochastic.solve_semi_dual_entropic(a, b, M, reg, + method, numItermax) + >>> print(asgd_pi) + + References + ---------- + + [Genevay et al., 2016] : + Stochastic Optimization for Large-scale Optimal Transport, + Advances in Neural Information Processing Systems (2016), + arXiv preprint arxiv:1605.08527. + ''' + + n_source = np.shape(M)[0] + alpha = np.zeros(n_source) + for i in range(n_source): + r = M[i, :] - beta + min_r = np.min(r) + exp_beta = np.exp(-(r - min_r) / reg) * b + alpha[i] = min_r - reg * np.log(np.sum(exp_beta)) + return alpha + + +def solve_semi_dual_entropic(a, b, M, reg, method, numItermax=10000, lr=None, + log=False): + ''' + Compute the transportation matrix to solve the regularized discrete + measures optimal transport max problem + + The function solves the following optimization problem: + + .. math:: + \gamma = arg\min_\gamma <\gamma,M>_F + reg\cdot\Omega(\gamma) + s.t. \gamma 1 = a + \gamma^T 1= b + \gamma \geq 0 + where : + - M is the (ns,nt) metric cost matrix + - :math:`\Omega` is the entropic regularization term + :math:`\Omega(\gamma)=\sum_{i,j} \gamma_{i,j}\log(\gamma_{i,j})` + - a and b are source and target weights (sum to 1) + The algorithm used for solving the problem is the SAG or ASGD algorithms + as proposed in [18]_ + + + Parameters + ---------- + + a : np.ndarray(ns,), + source measure + b : np.ndarray(nt,), + target measure + M : np.ndarray(ns, nt), + cost matrix + reg : float number, + Regularization term > 0 + methode : str, + used method (SAG or ASGD) + numItermax : int number + number of iteration + lr : float number + learning rate + n_source : int number + size of the source measure + n_target : int number + size of the target measure + log : bool, optional + record log if True + + Returns + ------- + + pi : np.ndarray(ns, nt) + transportation matrix + log : dict + log dictionary return only if log==True in parameters + + Examples + -------- + + >>> n_source = 7 + >>> n_target = 4 + >>> reg = 1 + >>> numItermax = 300000 + >>> a = ot.utils.unif(n_source) + >>> b = ot.utils.unif(n_target) + >>> rng = np.random.RandomState(0) + >>> X_source = rng.randn(n_source, 2) + >>> Y_target = rng.randn(n_target, 2) + >>> M = ot.dist(X_source, Y_target) + >>> method = "ASGD" + >>> asgd_pi = stochastic.solve_semi_dual_entropic(a, b, M, reg, + method, numItermax) + >>> print(asgd_pi) + + References + ---------- + + [Genevay et al., 2016] : + Stochastic Optimization for Large-scale Optimal Transport, + Advances in Neural Information Processing Systems (2016), + arXiv preprint arxiv:1605.08527. + ''' + + if method.lower() == "sag": + opt_beta = sag_entropic_transport(a, b, M, reg, numItermax, lr) + elif method.lower() == "asgd": + opt_beta = averaged_sgd_entropic_transport(a, b, M, reg, numItermax, lr) + else: + print("Please, select your method between SAG and ASGD") + return None + + opt_alpha = c_transform_entropic(b, M, reg, opt_beta) + pi = (np.exp((opt_alpha[:, None] + opt_beta[None, :] - M[:, :]) / reg) * + a[:, None] * b[None, :]) + + if log: + log = {} + log['alpha'] = opt_alpha + log['beta'] = opt_beta + return pi, log + else: + return pi + + +############################################################################## +# Optimization toolbox for DUAL problems +############################################################################## + + +def batch_grad_dual_alpha(M, reg, alpha, beta, batch_size, batch_alpha, + batch_beta): + ''' + Computes the partial gradient of F_\W_varepsilon + + Compute the partial gradient of the dual problem: + + ..math: + \forall i in batch_alpha, + grad_alpha_i = 1 * batch_size - + sum_{j in batch_beta} exp((alpha_i + beta_j - M_{i,j})/reg) + + where : + - M is the (ns,nt) metric cost matrix + - alpha, beta are dual variables in R^ixR^J + - reg is the regularization term + - batch_alpha and batch_beta are list of index + + The algorithm used for solving the dual problem is the SGD algorithm + as proposed in [19]_ [alg.1] + + Parameters + ---------- + + reg : float number, + Regularization term > 0 + M : np.ndarray(ns, nt), + cost matrix + alpha : np.ndarray(ns,) + dual variable + beta : np.ndarray(nt,) + dual variable + batch_size : int number + size of the batch + batch_alpha : np.ndarray(bs,) + batch of index of alpha + batch_beta : np.ndarray(bs,) + batch of index of beta + + Returns + ------- + + grad : np.ndarray(ns,) + partial grad F in alpha + + Examples + -------- + + >>> n_source = 7 + >>> n_target = 4 + >>> reg = 1 + >>> numItermax = 20000 + >>> lr = 0.1 + >>> batch_size = 3 + >>> log = True + >>> a = ot.utils.unif(n_source) + >>> b = ot.utils.unif(n_target) + >>> rng = np.random.RandomState(0) + >>> X_source = rng.randn(n_source, 2) + >>> Y_target = rng.randn(n_target, 2) + >>> M = ot.dist(X_source, Y_target) + >>> sgd_dual_pi, log = stochastic.solve_dual_entropic(a, b, M, reg, + batch_size, + numItermax, lr, log) + >>> print(log['alpha'], log['beta']) + >>> print(sgd_dual_pi) + + References + ---------- + + [Seguy et al., 2018] : + International Conference on Learning Representation (2018), + arXiv preprint arxiv:1711.02283. + ''' + + grad_alpha = np.zeros(batch_size) + grad_alpha[:] = batch_size + for j in batch_beta: + grad_alpha -= np.exp((alpha[batch_alpha] + beta[j] - + M[batch_alpha, j]) / reg) + return grad_alpha + + +def batch_grad_dual_beta(M, reg, alpha, beta, batch_size, batch_alpha, + batch_beta): + ''' + Computes the partial gradient of F_\W_varepsilon + + Compute the partial gradient of the dual problem: + + ..math: + \forall j in batch_beta, + grad_beta_j = 1 * batch_size - + sum_{i in batch_alpha} exp((alpha_i + beta_j - M_{i,j})/reg) + + where : + - M is the (ns,nt) metric cost matrix + - alpha, beta are dual variables in R^ixR^J + - reg is the regularization term + - batch_alpha and batch_beta are list of index + + The algorithm used for solving the dual problem is the SGD algorithm + as proposed in [19]_ [alg.1] + + Parameters + ---------- + + M : np.ndarray(ns, nt), + cost matrix + reg : float number, + Regularization term > 0 + alpha : np.ndarray(ns,) + dual variable + beta : np.ndarray(nt,) + dual variable + batch_size : int number + size of the batch + batch_alpha : np.ndarray(bs,) + batch of index of alpha + batch_beta : np.ndarray(bs,) + batch of index of beta + + Returns + ------- + + grad : np.ndarray(ns,) + partial grad F in beta + + Examples + -------- + + >>> n_source = 7 + >>> n_target = 4 + >>> reg = 1 + >>> numItermax = 20000 + >>> lr = 0.1 + >>> batch_size = 3 + >>> log = True + >>> a = ot.utils.unif(n_source) + >>> b = ot.utils.unif(n_target) + >>> rng = np.random.RandomState(0) + >>> X_source = rng.randn(n_source, 2) + >>> Y_target = rng.randn(n_target, 2) + >>> M = ot.dist(X_source, Y_target) + >>> sgd_dual_pi, log = stochastic.solve_dual_entropic(a, b, M, reg, + batch_size, + numItermax, lr, log) + >>> print(log['alpha'], log['beta']) + >>> print(sgd_dual_pi) + + References + ---------- + + [Seguy et al., 2018] : + International Conference on Learning Representation (2018), + arXiv preprint arxiv:1711.02283. + + ''' + + grad_beta = np.zeros(batch_size) + grad_beta[:] = batch_size + for i in batch_alpha: + grad_beta -= np.exp((alpha[i] + + beta[batch_beta] - M[i, batch_beta]) / reg) + return grad_beta + + +def sgd_entropic_regularization(M, reg, batch_size, numItermax, lr, + alternate=True): + ''' + Compute the sgd algorithm to solve the regularized discrete measures + optimal transport dual problem + + The function solves the following optimization problem: + + .. math:: + \gamma = arg\min_\gamma <\gamma,M>_F + reg\cdot\Omega(\gamma) + s.t. \gamma 1 = a + \gamma^T 1= b + \gamma \geq 0 + where : + - M is the (ns,nt) metric cost matrix + - :math:`\Omega` is the entropic regularization term + :math:`\Omega(\gamma)=\sum_{i,j} \gamma_{i,j}\log(\gamma_{i,j})` + - a and b are source and target weights (sum to 1) + + Parameters + ---------- + + M : np.ndarray(ns, nt), + cost matrix + reg : float number, + Regularization term > 0 + batch_size : int number + size of the batch + numItermax : int number + number of iteration + lr : float number + learning rate + alternate : bool, optional + alternating algorithm + + Returns + ------- + + alpha : np.ndarray(ns,) + dual variable + beta : np.ndarray(nt,) + dual variable + + Examples + -------- + + >>> n_source = 7 + >>> n_target = 4 + >>> reg = 1 + >>> numItermax = 20000 + >>> lr = 0.1 + >>> batch_size = 3 + >>> log = True + >>> a = ot.utils.unif(n_source) + >>> b = ot.utils.unif(n_target) + >>> rng = np.random.RandomState(0) + >>> X_source = rng.randn(n_source, 2) + >>> Y_target = rng.randn(n_target, 2) + >>> M = ot.dist(X_source, Y_target) + >>> sgd_dual_pi, log = stochastic.solve_dual_entropic(a, b, M, reg, + batch_size, + numItermax, lr, log) + >>> print(log['alpha'], log['beta']) + >>> print(sgd_dual_pi) + + References + ---------- + + [Seguy et al., 2018] : + International Conference on Learning Representation (2018), + arXiv preprint arxiv:1711.02283. + ''' + + n_source = np.shape(M)[0] + n_target = np.shape(M)[1] + cur_alpha = np.random.randn(n_source) + cur_beta = np.random.randn(n_target) + if alternate: + for cur_iter in range(numItermax): + k = np.sqrt(cur_iter + 1) + batch_alpha = np.random.choice(n_source, batch_size, replace=False) + batch_beta = np.random.choice(n_target, batch_size, replace=False) + grad_F_alpha = batch_grad_dual_alpha(M, reg, cur_alpha, cur_beta, + batch_size, batch_alpha, + batch_beta) + cur_alpha[batch_alpha] += (lr / k) * grad_F_alpha + grad_F_beta = batch_grad_dual_beta(M, reg, cur_alpha, cur_beta, + batch_size, batch_alpha, + batch_beta) + cur_beta[batch_beta] += (lr / k) * grad_F_beta + + else: + for cur_iter in range(numItermax): + k = np.sqrt(cur_iter + 1) + batch_alpha = np.random.choice(n_source, batch_size, replace=False) + batch_beta = np.random.choice(n_target, batch_size, replace=False) + grad_F_alpha = batch_grad_dual_alpha(M, reg, cur_alpha, cur_beta, + batch_size, batch_alpha, + batch_beta) + grad_F_beta = batch_grad_dual_beta(M, reg, cur_alpha, cur_beta, + batch_size, batch_alpha, + batch_beta) + cur_alpha[batch_alpha] += (lr / k) * grad_F_alpha + cur_beta[batch_beta] += (lr / k) * grad_F_beta + + return cur_alpha, cur_beta + + +def solve_dual_entropic(a, b, M, reg, batch_size, numItermax=10000, lr=1, + log=False): + ''' + Compute the transportation matrix to solve the regularized discrete measures + optimal transport dual problem + + The function solves the following optimization problem: + + .. math:: + \gamma = arg\min_\gamma <\gamma,M>_F + reg\cdot\Omega(\gamma) + s.t. \gamma 1 = a + \gamma^T 1= b + \gamma \geq 0 + where : + - M is the (ns,nt) metric cost matrix + - :math:`\Omega` is the entropic regularization term + :math:`\Omega(\gamma)=\sum_{i,j} \gamma_{i,j}\log(\gamma_{i,j})` + - a and b are source and target weights (sum to 1) + + Parameters + ---------- + + a : np.ndarray(ns,), + source measure + b : np.ndarray(nt,), + target measure + M : np.ndarray(ns, nt), + cost matrix + reg : float number, + Regularization term > 0 + batch_size : int number + size of the batch + numItermax : int number + number of iteration + lr : float number + learning rate + log : bool, optional + record log if True + + Returns + ------- + + pi : np.ndarray(ns, nt) + transportation matrix + log : dict + log dictionary return only if log==True in parameters + + Examples + -------- + + >>> n_source = 7 + >>> n_target = 4 + >>> reg = 1 + >>> numItermax = 20000 + >>> lr = 0.1 + >>> batch_size = 3 + >>> log = True + >>> a = ot.utils.unif(n_source) + >>> b = ot.utils.unif(n_target) + >>> rng = np.random.RandomState(0) + >>> X_source = rng.randn(n_source, 2) + >>> Y_target = rng.randn(n_target, 2) + >>> M = ot.dist(X_source, Y_target) + >>> sgd_dual_pi, log = stochastic.solve_dual_entropic(a, b, M, reg, + batch_size, + numItermax, lr, log) + >>> print(log['alpha'], log['beta']) + >>> print(sgd_dual_pi) + + References + ---------- + + [Seguy et al., 2018] : + International Conference on Learning Representation (2018), + arXiv preprint arxiv:1711.02283. + ''' + + opt_alpha, opt_beta = sgd_entropic_regularization(M, reg, batch_size, + numItermax, lr) + pi = (np.exp((opt_alpha[:, None] + opt_beta[None, :] - M[:, :]) / reg) * + a[:, None] * b[None, :]) + if log: + log = {} + log['alpha'] = opt_alpha + log['beta'] = opt_beta + return pi, log + else: + return pi diff --git a/test/test_stochastic.py b/test/test_stochastic.py new file mode 100644 index 0000000..f315c88 --- /dev/null +++ b/test/test_stochastic.py @@ -0,0 +1,191 @@ +""" +========================== +Stochastic test +========================== + +This example is designed to test the stochatic optimization algorithms module +for descrete and semicontinous measures from the POT library. + +""" + +# Author: Kilian Fatras <kilian.fatras@gmail.com> +# +# License: MIT License + +import numpy as np +import ot + + +############################################################################# +# COMPUTE TEST FOR SEMI-DUAL PROBLEM +############################################################################# + +############################################################################# +# +# TEST SAG algorithm +# --------------------------------------------- +# 2 identical discrete measures u defined on the same space with a +# regularization term, a learning rate and a number of iteration + + +def test_stochastic_sag(): + # test sag + n = 15 + reg = 1 + numItermax = 300000 + rng = np.random.RandomState(0) + + x = rng.randn(n, 2) + u = ot.utils.unif(n) + + M = ot.dist(x, x) + + G = ot.stochastic.solve_semi_dual_entropic(u, u, M, reg, "sag", + numItermax=numItermax) + + # check constratints + np.testing.assert_allclose( + u, G.sum(1), atol=1e-04) # cf convergence sag + np.testing.assert_allclose( + u, G.sum(0), atol=1e-04) # cf convergence sag + + +############################################################################# +# +# TEST ASGD algorithm +# --------------------------------------------- +# 2 identical discrete measures u defined on the same space with a +# regularization term, a learning rate and a number of iteration + + +def test_stochastic_asgd(): + # test asgd + n = 15 + reg = 1 + numItermax = 300000 + rng = np.random.RandomState(0) + + x = rng.randn(n, 2) + u = ot.utils.unif(n) + + M = ot.dist(x, x) + + G = ot.stochastic.solve_semi_dual_entropic(u, u, M, reg, "asgd", + numItermax=numItermax) + + # check constratints + np.testing.assert_allclose( + u, G.sum(1), atol=1e-03) # cf convergence asgd + np.testing.assert_allclose( + u, G.sum(0), atol=1e-03) # cf convergence asgd + + +############################################################################# +# +# TEST Convergence SAG and ASGD toward Sinkhorn's solution +# -------------------------------------------------------- +# 2 identical discrete measures u defined on the same space with a +# regularization term, a learning rate and a number of iteration + + +def test_sag_asgd_sinkhorn(): + # test all algorithms + n = 15 + reg = 1 + nb_iter = 300000 + rng = np.random.RandomState(0) + + x = rng.randn(n, 2) + u = ot.utils.unif(n) + zero = np.zeros(n) + M = ot.dist(x, x) + + G_asgd = ot.stochastic.solve_semi_dual_entropic(u, u, M, reg, "asgd", + numItermax=nb_iter) + G_sag = ot.stochastic.solve_semi_dual_entropic(u, u, M, reg, "sag", + numItermax=nb_iter) + G_sinkhorn = ot.sinkhorn(u, u, M, reg) + + # check constratints + np.testing.assert_allclose( + zero, (G_sag - G_sinkhorn).sum(1), atol=1e-03) # cf convergence sag + np.testing.assert_allclose( + zero, (G_sag - G_sinkhorn).sum(0), atol=1e-03) # cf convergence sag + np.testing.assert_allclose( + zero, (G_asgd - G_sinkhorn).sum(1), atol=1e-03) # cf convergence asgd + np.testing.assert_allclose( + zero, (G_asgd - G_sinkhorn).sum(0), atol=1e-03) # cf convergence asgd + np.testing.assert_allclose( + G_sag, G_sinkhorn, atol=1e-03) # cf convergence sag + np.testing.assert_allclose( + G_asgd, G_sinkhorn, atol=1e-03) # cf convergence asgd + + +############################################################################# +# COMPUTE TEST FOR DUAL PROBLEM +############################################################################# + +############################################################################# +# +# TEST SGD algorithm +# --------------------------------------------- +# 2 identical discrete measures u defined on the same space with a +# regularization term, a batch_size and a number of iteration + + +def test_stochastic_dual_sgd(): + # test sgd + n = 10 + reg = 1 + numItermax = 300000 + batch_size = 8 + rng = np.random.RandomState(0) + + x = rng.randn(n, 2) + u = ot.utils.unif(n) + + M = ot.dist(x, x) + + G = ot.stochastic.solve_dual_entropic(u, u, M, reg, batch_size, + numItermax=numItermax) + + # check constratints + np.testing.assert_allclose( + u, G.sum(1), atol=1e-02) # cf convergence sgd + np.testing.assert_allclose( + u, G.sum(0), atol=1e-02) # cf convergence sgd + + +############################################################################# +# +# TEST Convergence SGD toward Sinkhorn's solution +# -------------------------------------------------------- +# 2 identical discrete measures u defined on the same space with a +# regularization term, a batch_size and a number of iteration + + +def test_dual_sgd_sinkhorn(): + # test all dual algorithms + n = 10 + reg = 1 + nb_iter = 300000 + batch_size = 8 + rng = np.random.RandomState(0) + + x = rng.randn(n, 2) + u = ot.utils.unif(n) + zero = np.zeros(n) + M = ot.dist(x, x) + + G_sgd = ot.stochastic.solve_dual_entropic(u, u, M, reg, batch_size, + numItermax=nb_iter) + + G_sinkhorn = ot.sinkhorn(u, u, M, reg) + + # check constratints + np.testing.assert_allclose( + zero, (G_sgd - G_sinkhorn).sum(1), atol=1e-02) # cf convergence sgd + np.testing.assert_allclose( + zero, (G_sgd - G_sinkhorn).sum(0), atol=1e-02) # cf convergence sgd + np.testing.assert_allclose( + G_sgd, G_sinkhorn, atol=1e-02) # cf convergence sgd |