Version 1.5.1 patch release (#114)

* Add support for tensorflow backend which allows for differentiability (#112) * Added support for tensorflow * Updates to get tests passing * Or --> And * Moving modopt to allow working with tensorflow * Fix issues with wos * Fix all flakes finally! * Update modopt/base/backend.py Co-authored-by: Samuel Farrens <samuel.farrens@gmail.com> * Update modopt/base/backend.py Co-authored-by: Samuel Farrens <samuel.farrens@gmail.com> * Minute updates to codes * Add dynamic module * Fix docu * Fix PEP Co-authored-by: chaithyagr <chaithyagr@gitlab.com> Co-authored-by: Samuel Farrens <samuel.farrens@gmail.com> * Fix 115 (#116) * Fix issues * Add right tests * Fix PEP Co-authored-by: chaithyagr <chaithyagr@gitlab.com> * Minor bug fix, remove elif (#124) Co-authored-by: chaithyagr <chaithyagr@gitlab.com> * Add tests for modopt.base.backend and fix minute bug uncovered (#126) * Minor bug fix, remove elif * Add tests for backend * Fix tests * Add tests * Remove cupy * PEP fixes * Fix PEP * Fix PEP and update * Final PEP * Update setup.cfg Co-authored-by: Samuel Farrens <samuel.farrens@gmail.com> * Update test_base.py Co-authored-by: chaithyagr <chaithyagr@gitlab.com> Co-authored-by: Samuel Farrens <samuel.farrens@gmail.com> * Release cleanup (#128) * updated GPU dependencies * added logo to manifest * updated package version and release date Co-authored-by: Chaithya G R <chaithyagr@gmail.com> Co-authored-by: chaithyagr <chaithyagr@gitlab.com>
CEA-COSMIC · Apr 22, 2021 · af01498 · af01498
1 parent 616c9e7
commit af01498
Show file tree

Hide file tree

Showing 12 changed files with 237 additions and 82 deletions.
diff --git a/.github/workflows/ci-build.yml b/.github/workflows/ci-build.yml
@@ -53,6 +53,7 @@ jobs:
           python -m pip install -r develop.txt
           python -m pip install -r docs/requirements.txt
           python -m pip install astropy scikit-image scikit-learn
+          python -m pip install tensorflow>=2.4.1
           python -m pip install twine
           python -m pip install .
 

diff --git a/MANIFEST.in b/MANIFEST.in
@@ -3,3 +3,4 @@ include develop.txt
 include docs/requirements.txt
 include README.rst
 include LICENSE.txt
+include docs/source/modopt_logo.png
diff --git a/README.md b/README.md
@@ -58,6 +58,7 @@ Note that none of these are required for running on a CPU.
 
 * [CuPy](https://cupy.dev/)
 * [Torch](https://pytorch.org/)
+* [TensorFlow](https://www.tensorflow.org/)
 
 ## Citation
 

diff --git a/docs/source/dependencies.rst b/docs/source/dependencies.rst
@@ -82,6 +82,7 @@ For GPU compliance the following packages can also be installed:
 
 * |link-to-cupy|
 * |link-to-torch|
+* |link-to-tf|
 
 .. |link-to-cupy| raw:: html
 
@@ -93,6 +94,11 @@ For GPU compliance the following packages can also be installed:
   <a href="https://pytorch.org/"
   target="_blank">Torch</a>
 
+.. |link-to-tf| raw:: html
+
+  <a href="https://www.tensorflow.org/"
+  target="_blank">TensorFlow</a>
+
 .. note::
 
   Note that none of these are required for running on a CPU.
diff --git a/docs/source/index.rst b/docs/source/index.rst
@@ -12,8 +12,8 @@ ModOpt Documentation
 .. include:: toc.rst
 
 :Author: Samuel Farrens `(samuel.farrens@cea.fr) <samuel.farrens@cea.fr>`_
-:Version: 1.5.0
-:Release Date: 31/03/2021
+:Version: 1.5.1
+:Release Date: 22/04/2021
 :Repository: |link-to-repo|
 
 .. |link-to-repo| raw:: html

diff --git a/modopt/base/backend.py b/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']
+    if 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()
diff --git a/modopt/math/matrix.py b/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()
 

diff --git a/modopt/opt/algorithms.py b/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()
@@ -789,7 +777,7 @@ def _update(self):
         if self._cost_func:
             self.converge = (
                 self.any_convergence_flag()
-                or self._cost_func.get_cost(self._x_new),
+                or self._cost_func.get_cost(self._x_new)
             )
 
     def iterate(self, max_iter=150):
@@ -1548,7 +1536,7 @@ def _update(self):
         if self._cost_func:
             self.converge = (
                 self.any_convergence_flag()
-                or self._cost_func.get_cost(self._x_new),
+                or self._cost_func.get_cost(self._x_new)
             )
 
     def iterate(self, max_iter=150):