diff --git a/modopt/base/backend.py b/modopt/base/backend.py index 1f5a7a3a..def77151 100644 --- 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'] + 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() diff --git a/modopt/math/matrix.py b/modopt/math/matrix.py index cb54cebc..be737f52 100644 --- 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,7 +298,7 @@ def __init__( data_shape, data_type=float, auto_run=True, - use_gpu=False, + compute_backend='numpy', verbose=False, ): @@ -311,10 +306,9 @@ def __init__( 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 index ccd8fe21..fea737cb 100644 --- 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() diff --git a/setup.cfg b/setup.cfg index 56a40d12..74fb3f79 100644 --- a/setup.cfg +++ b/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