From baf51b5998ce1b43d22b61c041146c0e995ffc3f Mon Sep 17 00:00:00 2001 From: Nolwen Date: Fri, 1 Dec 2023 17:02:53 +0100 Subject: [PATCH] Remove direct tensorflow imports from tests To be fully agnostic, we prefer having utility functions that are checking the backend use like now `is_in_gpu_mode()`. --- tests/conftest.py | 28 ++++++++++++++++++++++++---- tests/test_backward_conv.py | 3 +-- tests/test_backward_layers.py | 3 +-- tests/test_backward_native_layers.py | 3 +-- tests/test_conv.py | 3 +-- tests/test_utils_conv.py | 4 ++-- 6 files changed, 30 insertions(+), 14 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index d995a423..14c3967e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -21,6 +21,12 @@ from numpy.testing import assert_almost_equal from decomon.core import ForwardMode, Slope +from decomon.keras_utils import ( + BACKEND_JAX, + BACKEND_NUMPY, + BACKEND_PYTORCH, + BACKEND_TENSORFLOW, +) from decomon.models.utils import ConvertMethod from decomon.types import Tensor @@ -165,10 +171,24 @@ def __call__(self, inputs_: List[np.ndarray]): class Helpers: @staticmethod - def tensorflow_in_GPU_mode() -> bool: - import tensorflow - - return len(tensorflow.config.list_physical_devices("GPU")) > 0 + def in_GPU_mode() -> bool: + backend = keras.config.backend() + if backend == BACKEND_TENSORFLOW: + import tensorflow + + return len(tensorflow.config.list_physical_devices("GPU")) > 0 + elif backend == BACKEND_PYTORCH: + import torch + + return torch.cuda.is_available() + elif backend == BACKEND_NUMPY: + return False + elif backend == BACKEND_JAX: + import jax + + return jax.devices()[0].platform != "cpu" + else: + raise NotImplementedError(f"Not implemented for {backend} backend.") @staticmethod def is_method_mode_compatible(method, mode): diff --git a/tests/test_backward_conv.py b/tests/test_backward_conv.py index 05bcc5bb..788b84f1 100644 --- a/tests/test_backward_conv.py +++ b/tests/test_backward_conv.py @@ -4,14 +4,13 @@ import keras.config as keras_config import numpy as np import pytest -from tensorflow.python.keras.backend import _get_available_gpus from decomon.backward_layers.convert import to_backward from decomon.layers.decomon_layers import DecomonConv2D def test_Decomon_conv_box(data_format, padding, use_bias, mode, floatx, decimal, helpers): - if data_format == "channels_first" and not helpers.tensorflow_in_GPU_mode(): + if data_format == "channels_first" and not helpers.in_GPU_mode(): pytest.skip("data format 'channels first' is possible only in GPU mode") odd, m_0, m_1 = 0, 0, 1 diff --git a/tests/test_backward_layers.py b/tests/test_backward_layers.py index cd3abb27..ad58e2ab 100644 --- a/tests/test_backward_layers.py +++ b/tests/test_backward_layers.py @@ -1,7 +1,6 @@ import keras.config as keras_config import pytest from keras.layers import Layer, Reshape -from tensorflow.python.keras.backend import _get_available_gpus from decomon.backward_layers.convert import to_backward from decomon.core import ForwardMode, Slope @@ -105,7 +104,7 @@ def test_Backward_Activation_multiD_box(odd, activation, floatx, decimal, mode, def test_Backward_Flatten_multiD_box(odd, floatx, decimal, mode, data_format, helpers): - if data_format == "channels_first" and not helpers.tensorflow_in_GPU_mode(): + if data_format == "channels_first" and not helpers.in_GPU_mode(): pytest.skip("data format 'channels first' is possible only in GPU mode") dc_decomp = False diff --git a/tests/test_backward_native_layers.py b/tests/test_backward_native_layers.py index 9fc68737..838b24f5 100644 --- a/tests/test_backward_native_layers.py +++ b/tests/test_backward_native_layers.py @@ -1,7 +1,6 @@ import keras.config as keras_config import pytest from keras.layers import Activation, Flatten, Reshape -from tensorflow.python.keras.backend import _get_available_gpus from decomon.backward_layers.convert import to_backward @@ -65,7 +64,7 @@ def test_Backward_NativeActivation_multiD_box(odd, activation, floatx, decimal, def test_Backward_NativeFlatten_multiD_box(odd, floatx, decimal, mode, data_format, helpers): - if data_format == "channels_first" and not helpers.tensorflow_in_GPU_mode(): + if data_format == "channels_first" and not helpers.in_GPU_mode(): pytest.skip("data format 'channels first' is possible only in GPU mode") dc_decomp = False diff --git a/tests/test_conv.py b/tests/test_conv.py index 2ea86843..5f98ccbb 100644 --- a/tests/test_conv.py +++ b/tests/test_conv.py @@ -5,7 +5,6 @@ import numpy as np import pytest from keras.layers import Conv2D -from tensorflow.python.keras.backend import _get_available_gpus from decomon.core import ForwardMode, get_affine, get_ibp from decomon.layers.convert import to_decomon @@ -13,7 +12,7 @@ def test_Decomon_conv_box(data_format, mode, dc_decomp, floatx, decimal, helpers): - if data_format == "channels_first" and not helpers.tensorflow_in_GPU_mode(): + if data_format == "channels_first" and not helpers.in_GPU_mode(): pytest.skip("data format 'channels first' is possible only in GPU mode") odd, m_0, m_1 = 0, 0, 1 diff --git a/tests/test_utils_conv.py b/tests/test_utils_conv.py index 0b25ff2e..0f0979d3 100644 --- a/tests/test_utils_conv.py +++ b/tests/test_utils_conv.py @@ -17,7 +17,7 @@ def test_toeplitz_from_Keras(channels, filter_size, strides, flatten, data_forma if floatx == 16: decimal = 0 - if data_format == "channels_first" and not helpers.tensorflow_in_GPU_mode(): + if data_format == "channels_first" and not helpers.in_GPU_mode(): pytest.skip("data format 'channels first' is possible only in GPU mode") dc_decomp = False @@ -67,7 +67,7 @@ def test_toeplitz_from_Keras(channels, filter_size, strides, flatten, data_forma def test_toeplitz_from_Decomon( floatx, decimal, mode, channels, filter_size, strides, flatten, data_format, padding, helpers ): - if data_format == "channels_first" and not helpers.tensorflow_in_GPU_mode(): + if data_format == "channels_first" and not helpers.in_GPU_mode(): pytest.skip("data format 'channels first' is possible only in GPU mode") odd, m_0, m_1 = 0, 0, 1