[MRG] Debug convolutional methods that compute barycenters to work with different devices. (#533)

* [FEAT] Add the parameter 'type_as' to the backends * [TEST] add tests for the 'type_as' backends * [DEBUG] Debug dtype in pytorch * [FIX] Add type_as every time linspace is called * [TEST] Add test for the convolutional_barycenter2d algorithms (they are the only ones use linspace) * [DEBUG] PEP 8 * [DOC] Add the new changes to RELEASES.md * [REFACTOR] Minor refactor that checks the GPU on the last line * [DEBUG] pep8 * [REFACTOR] Add a function to generalize the creation of random images * [REFACTOR] Mantain th same style as before * Update gitignore --------- Co-authored-by: Francisco Muñoz <fmunoz@ug.uchile.cl>

Author: framunoz

.gitignore

@@ -97,6 +97,7 @@ celerybeat-schedule
 # virtualenv
 venv/
 ENV/
+.venv/
 
 # Spyder project settings
 .spyderproject
@@ -120,4 +121,4 @@ debug
 .vscode
 
 # pytest cahche
-.pytest_cache
+.pytest_cache

RELEASES.md

@@ -6,6 +6,8 @@
 + Added support for [Nearest Brenier Potentials (SSNB)](http://proceedings.mlr.press/v108/paty20a/paty20a.pdf) (PR #526)
 + Tweaked `get_backend` to ignore `None` inputs (PR #525)
 + Callbacks for generalized conditional gradient in `ot.da.sinkhorn_l1l2_gl` are now vectorized to improve performance (PR #507)
++ The `linspace` method of the backends now has the `type_as` argument to convert to the same dtype and device. (PR #533)
++ The `convolutional_barycenter2d` and `convolutional_barycenter2d_debiased` functions now work with different devices.. (PR #533)
 
 #### Closed issues
 - Fix line search evaluating cost outside of the interpolation range (Issue #502, PR #504)

ot/backend.py

@@ -86,15 +86,15 @@
 #
 # License: MIT License
 
-import numpy as np
 import os
-import scipy
-import scipy.linalg
-from scipy.sparse import issparse, coo_matrix, csr_matrix
-import scipy.special as special
 import time
 import warnings
 
+import numpy as np
+import scipy
+import scipy.linalg
+import scipy.special as special
+from scipy.sparse import coo_matrix, csr_matrix, issparse
 
 DISABLE_TORCH_KEY = 'POT_BACKEND_DISABLE_PYTORCH'
 DISABLE_JAX_KEY = 'POT_BACKEND_DISABLE_JAX'
@@ -650,7 +650,7 @@ def std(self, a, axis=None):
  """
  raise NotImplementedError()
 
- def linspace(self, start, stop, num):
+ def linspace(self, start, stop, num, type_as=None):
  r"""
  Returns a specified number of evenly spaced values over a given interval.
 
@@ -1208,8 +1208,11 @@ def median(self, a, axis=None):
  def std(self, a, axis=None):
  return np.std(a, axis=axis)
 
- def linspace(self, start, stop, num):
- return np.linspace(start, stop, num)
+ def linspace(self, start, stop, num, type_as=None):
+ if type_as is None:
+ return np.linspace(start, stop, num)
+ else:
+ return np.linspace(start, stop, num, dtype=type_as.dtype)
 
  def meshgrid(self, a, b):
  return np.meshgrid(a, b)
@@ -1579,8 +1582,11 @@ def median(self, a, axis=None):
  def std(self, a, axis=None):
  return jnp.std(a, axis=axis)
 
- def linspace(self, start, stop, num):
- return jnp.linspace(start, stop, num)
+ def linspace(self, start, stop, num, type_as=None):
+ if type_as is None:
+ return jnp.linspace(start, stop, num)
+ else:
+ return self._change_device(jnp.linspace(start, stop, num, dtype=type_as.dtype), type_as)
 
  def meshgrid(self, a, b):
  return jnp.meshgrid(a, b)
@@ -1986,6 +1992,7 @@ def concatenate(self, arrays, axis=0):
 
  def zero_pad(self, a, pad_width, value=0):
  from torch.nn.functional import pad
+
  # pad_width is an array of ndim tuples indicating how many 0 before and after
  # we need to add. We first need to make it compliant with torch syntax, that
  # starts with the last dim, then second last, etc.
@@ -2006,6 +2013,7 @@ def mean(self, a, axis=None):
 
  def median(self, a, axis=None):
  from packaging import version
+
  # Since version 1.11.0, interpolation is available
  if version.parse(torch.__version__) >= version.parse("1.11.0"):
  if axis is not None:
@@ -2026,8 +2034,11 @@ def std(self, a, axis=None):
  else:
  return torch.std(a, unbiased=False)
 
- def linspace(self, start, stop, num):
- return torch.linspace(start, stop, num, dtype=torch.float64)
+ def linspace(self, start, stop, num, type_as=None):
+ if type_as is None:
+ return torch.linspace(start, stop, num)
+ else:
+ return torch.linspace(start, stop, num, dtype=type_as.dtype, device=type_as.device)
 
  def meshgrid(self, a, b):
  try:
@@ -2427,8 +2438,12 @@ def median(self, a, axis=None):
  def std(self, a, axis=None):
  return cp.std(a, axis=axis)
 
- def linspace(self, start, stop, num):
- return cp.linspace(start, stop, num)
+ def linspace(self, start, stop, num, type_as=None):
+ if type_as is None:
+ return cp.linspace(start, stop, num)
+ else:
+ with cp.cuda.Device(type_as.device):
+ return cp.linspace(start, stop, num, dtype=type_as.dtype)
 
  def meshgrid(self, a, b):
  return cp.meshgrid(a, b)
@@ -2834,8 +2849,11 @@ def median(self, a, axis=None):
  def std(self, a, axis=None):
  return tnp.std(a, axis=axis)
 
- def linspace(self, start, stop, num):
- return tnp.linspace(start, stop, num)
+ def linspace(self, start, stop, num, type_as=None):
+ if type_as is None:
+ return tnp.linspace(start, stop, num)
+ else:
+ return tnp.linspace(start, stop, num, dtype=type_as.dtype)
 
  def meshgrid(self, a, b):
  return tnp.meshgrid(a, b)

ot/bregman.py

@@ -20,7 +20,8 @@
 import numpy as np
 from scipy.optimize import fmin_l_bfgs_b
 
-from ot.utils import unif, dist, list_to_array
+from ot.utils import dist, list_to_array, unif
+
 from .backend import get_backend
 
 
@@ -2217,11 +2218,11 @@ def _convolutional_barycenter2d(A, reg, weights=None, numItermax=10000,
 
  # build the convolution operator
  # this is equivalent to blurring on horizontal then vertical directions
- t = nx.linspace(0, 1, A.shape[1])
+ t = nx.linspace(0, 1, A.shape[1], type_as=A)
  [Y, X] = nx.meshgrid(t, t)
  K1 = nx.exp(-(X - Y) ** 2 / reg)
 
- t = nx.linspace(0, 1, A.shape[2])
+ t = nx.linspace(0, 1, A.shape[2], type_as=A)
  [Y, X] = nx.meshgrid(t, t)
  K2 = nx.exp(-(X - Y) ** 2 / reg)
 
@@ -2295,11 +2296,11 @@ def _convolutional_barycenter2d_log(A, reg, weights=None, numItermax=10000,
  err = 1
  # build the convolution operator
  # this is equivalent to blurring on horizontal then vertical directions
- t = nx.linspace(0, 1, width)
+ t = nx.linspace(0, 1, width, type_as=A)
  [Y, X] = nx.meshgrid(t, t)
  M1 = - (X - Y) ** 2 / reg
 
- t = nx.linspace(0, 1, height)
+ t = nx.linspace(0, 1, height, type_as=A)
  [Y, X] = nx.meshgrid(t, t)
  M2 = - (X - Y) ** 2 / reg
 
@@ -2452,11 +2453,11 @@ def _convolutional_barycenter2d_debiased(A, reg, weights=None, numItermax=10000,
 
  # build the convolution operator
  # this is equivalent to blurring on horizontal then vertical directions
- t = nx.linspace(0, 1, width)
+ t = nx.linspace(0, 1, width, type_as=A)
  [Y, X] = nx.meshgrid(t, t)
  K1 = nx.exp(-(X - Y) ** 2 / reg)
 
- t = nx.linspace(0, 1, height)
+ t = nx.linspace(0, 1, height, type_as=A)
  [Y, X] = nx.meshgrid(t, t)
  K2 = nx.exp(-(X - Y) ** 2 / reg)
 
@@ -2532,11 +2533,11 @@ def _convolutional_barycenter2d_debiased_log(A, reg, weights=None, numItermax=10
  err = 1
  # build the convolution operator
  # this is equivalent to blurring on horizontal then vertical directions
- t = nx.linspace(0, 1, width)
+ t = nx.linspace(0, 1, width, type_as=A)
  [Y, X] = nx.meshgrid(t, t)
  M1 = - (X - Y) ** 2 / reg
 
- t = nx.linspace(0, 1, height)
+ t = nx.linspace(0, 1, height, type_as=A)
  [Y, X] = nx.meshgrid(t, t)
  M2 = - (X - Y) ** 2 / reg
 

test/test_backend.py

@@ -6,16 +6,13 @@
 #
 # License: MIT License
 
-import ot
-import ot.backend
-from ot.backend import torch, jax, tf
-
-import pytest
-
 import numpy as np
+import pytest
 from numpy.testing import assert_array_almost_equal_nulp
 
-from ot.backend import get_backend, get_backend_list, to_numpy
+import ot
+import ot.backend
+from ot.backend import get_backend, get_backend_list, jax, tf, to_numpy, torch
 
 
 def test_get_backend_list():
@@ -507,6 +504,7 @@ def test_func_backends(nx):
  lst_name.append('std')
 
  A = nx.linspace(0, 1, 50)
+ A = nx.linspace(0, 1, 50, type_as=Mb)
  lst_b.append(nx.to_numpy(A))
  lst_name.append('linspace')