Add support for tensorflow backend which allows for differentiability

modopt/base/backend.py

@@ -8,35 +8,73 @@
 
 """
 
-import warnings
 from importlib import util
 
 import numpy as np
 
+from modopt.interface.errors import warn
+
 try:
     import torch
+    from torch.utils.dlpack import from_dlpack as torch_from_dlpack
+    from torch.utils.dlpack import to_dlpack as torch_to_dlpack
+
 except ImportError:  # pragma: no cover
     import_torch = False
 else:
     import_torch = True
 
 # Handle the compatibility with variable
-gpu_compatibility = {
-    'cupy': False,
-    'cupy-cudnn': False,
+LIBRARIES = {
+    'cupy': None,
+    'tensorflow': None,
+    'numpy': np,
 }
 
 if util.find_spec('cupy') is not None:
     try:
         import cupy as cp
-        gpu_compatibility['cupy'] = True
+        LIBRARIES['cupy'] = cp
+    except ImportError:
+        pass
 
-        if util.find_spec('cupy.cuda.cudnn') is not None:
-            gpu_compatibility['cupy-cudnn'] = True
+if util.find_spec('tensorflow') is not None:
+    try:
+        from tensorflow.experimental import numpy as tnp
+        LIBRARIES['tensorflow'] = tnp
     except ImportError:
         pass
 
 
+def get_backend(backend):
+    """Get backend.
+
+    Returns the backend module for input specified by string
+
+    Parameters
+    ----------
+    backend: str
+        String holding the backend name. One of `tensorflow`,
+        `numpy` or `cupy`.
+
+    Returns
+    -------
+    tuple
+        Returns the module for carrying out calculations and the actual backend
+        that was reverted towards. If the right libraries are not installed,
+        the function warns and reverts to `numpy` backend
+    """
+    if backend not in LIBRARIES.keys() or LIBRARIES[backend] is None:
+        msg = (
+            '{0} backend not possible, please ensure that '
+            + 'the optional libraries are installed.\n'
+            + 'Reverting to numpy'
+        )
+        warn(msg.format(backend))
+        backend = 'numpy'
+    return LIBRARIES[backend], backend
+
+
 def get_array_module(input_data):
     """Get Array Module.
 
@@ -54,48 +92,47 @@ def get_array_module(input_data):
         The numpy or cupy module
 
     """
-    if gpu_compatibility['cupy']:
-        return cp.get_array_module(input_data)
-
+    if LIBRARIES['tensorflow'] is not None:
+        if isinstance(input_data, LIBRARIES['tensorflow'].ndarray):
+            return LIBRARIES['tensorflow']
+    elif LIBRARIES['cupy'] is not None:
+        if isinstance(input_data, LIBRARIES['cupy'].ndarray):
+            return LIBRARIES['cupy']
     return np
 
 
-def move_to_device(input_data):
+def change_backend(input_data, backend='cupy'):
     """Move data to device.
 
-    This method moves data from CPU to GPU if we have the
-    compatibility to do so. It returns the same data if
-    it is already on GPU.
+    This method changes the backend of an array
+    This can be used to copy data to GPU or to CPU
 
     Parameters
     ----------
     input_data : numpy.ndarray or cupy.ndarray
         Input data array to be moved
+    backend: str, optional
+        The backend to use, one among `tensorflow`, `cupy` and
+        `numpy`. Default is `cupy`.
 
     Returns
     -------
-    cupy.ndarray
-        The CuPy array residing on GPU
+    backend.ndarray
+        An ndarray of specified backend
 
     """
     xp = get_array_module(input_data)
-
-    if xp == cp:
+    txp, target_backend = get_backend(backend)
+    if xp == txp:
         return input_data
-
-    if gpu_compatibility['cupy']:
-        return cp.array(input_data)
-
-    warnings.warn('Cupy is not installed, cannot move data to GPU')
-
-    return input_data
+    return txp.array(input_data)
 
 
 def move_to_cpu(input_data):
     """Move data to CPU.
 
-    This method moves data from GPU to CPU.It returns the same data if it is
-    already on CPU.
+    This method moves data from GPU to CPU.
+    It returns the same data if it is already on CPU.
 
     Parameters
     ----------
@@ -107,13 +144,20 @@ def move_to_cpu(input_data):
     numpy.ndarray
         The NumPy array residing on CPU
 
+    Raises
+    ------
+    ValueError
+        if the input does not correspond to any array
     """
     xp = get_array_module(input_data)
 
-    if xp == np:
+    if xp == LIBRARIES['numpy']:
         return input_data
-
-    return input_data.get()
+    elif xp == LIBRARIES['cupy']:
+        return input_data.get()
+    elif xp == LIBRARIES['tensorflow']:
+        return input_data.data.numpy()
+    raise ValueError('Cannot identify the array type.')
 
 
 def convert_to_tensor(input_data):
@@ -150,7 +194,7 @@ def convert_to_tensor(input_data):
     if xp == np:
         return torch.Tensor(input_data)
 
-    return torch.utils.dlpack.from_dlpack(input_data.toDlpack()).float()
+    return torch_from_dlpack(input_data.toDlpack()).float()
 
 
 def convert_to_cupy_array(input_data):
@@ -182,6 +226,6 @@ def convert_to_cupy_array(input_data):
         )
 
     if input_data.is_cuda:
-        return cp.fromDlpack(torch.utils.dlpack.to_dlpack(input_data))
+        return cp.fromDlpack(torch_to_dlpack(input_data))
 
     return input_data.detach().numpy()

modopt/math/matrix.py

@@ -12,12 +12,7 @@
 
 import numpy as np
 
-from modopt.base.backend import get_array_module
-
-try:
-    import cupy as cp
-except ImportError:  # pragma: no cover
-    pass
+from modopt.base.backend import get_array_module, get_backend
 
 
 def gram_schmidt(matrix, return_opt='orthonormal'):
@@ -303,18 +298,17 @@ def __init__(
         data_shape,
         data_type=float,
         auto_run=True,
-        use_gpu=False,
+        compute_backend='numpy',
         verbose=False,
     ):
 
         self._operator = operator
         self._data_shape = data_shape
         self._data_type = data_type
         self._verbose = verbose
-        if use_gpu:
-            self.xp = cp
-        else:
-            self.xp = np
+        xp, compute_backend = get_backend(compute_backend)
+        self.xp = xp
+        self.compute_backend = compute_backend
         if auto_run:
             self.get_spec_rad()
 

modopt/opt/algorithms.py

@@ -54,11 +54,6 @@
 from modopt.opt.cost import costObj
 from modopt.opt.linear import Identity
 
-try:
-    import cupy as cp
-except ImportError:  # pragma: no cover
-    pass
-
 
 class SetUp(Observable):
     r"""Algorithm Set-Up.
@@ -92,7 +87,7 @@ def __init__(
         verbose=False,
         progress=True,
         step_size=None,
-        use_gpu=False,
+        compute_backend='numpy',
         **dummy_kwargs,
     ):
 
@@ -123,20 +118,9 @@ def __init__(
             )
             self.add_observer('cv_metrics', observer)
 
-        # Check for GPU
-        if use_gpu:
-            if backend.gpu_compatibility['cupy']:
-                self.xp = cp
-            else:
-                warn(
-                    'CuPy is not installed, cannot run on GPU!'
-                    + 'Running optimization on CPU.',
-                )
-                self.xp = np
-                use_gpu = False
-        else:
-            self.xp = np
-        self.use_gpu = use_gpu
+        xp, compute_backend = backend.get_backend(compute_backend)
+        self.xp = xp
+        self.compute_backend = compute_backend
 
     @property
     def metrics(self):
@@ -148,7 +132,9 @@ def metrics(self, metrics):
 
         if isinstance(metrics, type(None)):
             self._metrics = {}
-        elif not isinstance(metrics, dict):
+        elif isinstance(metrics, dict):
+            self._metrics = metrics
+        else:
             raise TypeError(
                 'Metrics must be a dictionary, not {0}.'.format(type(metrics)),
             )
@@ -184,10 +170,10 @@ def copy_data(self, input_data):
             Copy of input data
 
         """
-        if self.use_gpu:
-            return backend.move_to_device(input_data)
-
-        return self.xp.copy(input_data)
+        return self.xp.copy(backend.change_backend(
+            input_data,
+            self.compute_backend,
+        ))
 
     def _check_input_data(self, input_data):
         """Check input data type.
@@ -205,8 +191,10 @@ def _check_input_data(self, input_data):
             For invalid input type
 
         """
-        if not isinstance(input_data, self.xp.ndarray):
-            raise TypeError('Input data must be a numpy array.')
+        if not (isinstance(input_data, (self.xp.ndarray, np.ndarray))):
+            raise TypeError(
+                'Input data must be a numpy array or backend array',
+            )
 
     def _check_param(self, param_val):
         """Check algorithm parameters.
@@ -779,8 +767,8 @@ def _update(self):
             self._z_new = self._x_new
 
         # Update old values for next iteration.
-        self.xp.copyto(self._x_old, self._x_new)
-        self.xp.copyto(self._z_old, self._z_new)
+        self._x_old = self.xp.copy(self._x_new)
+        self._z_old = self.xp.copy(self._z_new)
 
         # Update parameter values for next iteration.
         self._update_param()

setup.cfg

@@ -26,7 +26,8 @@ per-file-ignores =
   #Justification: Needed for keeping package version and current API
   *__init__.py*: F401,F403,WPS347,WPS410,WPS412
   #Todo: Rethink conditional imports
-  modopt/base/backend.py: WPS229, WPS420
+  #Todo: How can we bypass mutable constants?
+  modopt/base/backend.py: WPS229, WPS420, WPS407
   #Todo: Rethink conditional imports
   modopt/base/observable.py: WPS420,WPS604
   #Todo: Check string for log formatting