Device selection with pytest. (#1713)

* The device to execute each test is set in stone.
* Update CONTRIBUTING.md
* Removed device placement logging and added comment.

Author: gabrieldemarmiesse

CONTRIBUTING.md

@@ -297,6 +297,10 @@ Run selected tests:
 python3 -m pytest path/to/file/or/directory/to/test
 ```
 
+Run the gpu only tests with `pytest -m needs_gpu ./tensorflow_addons`.
+Run the cpu only tests with `pytest -m 'not needs_gpu' ./tensorflow_addons`.
+
+
 #### Testing with Bazel
 
 Testing with Bazel is still supported but not recommended unless you have prior experience 
@@ -411,22 +415,46 @@ on Tensors, `if` or `for` for example. Or with `TensorArray`. In short, when the
  conversion to graph is not trivial. No need to use it on all
 your tests. Having fast tests is important.
 
-#### cpu_and_gpu
+#### Selecting the devices to run the test
+
+By default, each test is wrapped behind the scenes with a 
+```python
+with tf.device("CPU:0"):
+    ...
+```
 
-Will run your test function twice, once with `with tf.device("/device:CPU:0")` and 
-once with `with tf.device("/device:GPU:0")`. If a GPU is not present on the system, 
-the second test is skipped. To use it:
+This is automatic. But it's also possible to ask the test runner to run 
+the test twice, on CPU and on GPU, or only on GPU. Here is how to do it.
 
 ```python
-@pytest.mark.usefixtures("cpu_and_gpu")
+import pytest
+
+@pytest.mark.with_device(["cpu", "gpu"])
 def test_something():
-    assert ...== ...
+    # the code here will run twice, once on gpu, once on cpu.
+    ...
+
+@pytest.mark.with_device(["gpu"])
+def test_something_else():
+    # This test will be only run on gpu.
+    # The test runner will call with tf.device("GPU:0") behind the scenes.  
+    ...
+
+@pytest.mark.with_device(["cpu"])
+def test_something_more():
+    # Don't do that, this is the default behavior. 
+    ...
 ```
 
+Note that if a gpu is not detected on the system, the test will be 
+skipped and not marked as failed. Only the first gpu of the system is used,
+even when running pytest in multiprocessing mode. (`-n` argument). 
+Beware of the out of cuda memory errors if the number of pytest workers is too high.
+
 ##### When to use it?
 
-When you test custom CUDA code. We can expect existing TensorFlow ops to behave the same 
-on CPU and GPU.
+When you test custom CUDA code or float16 ops.
+We can expect other existing TensorFlow ops to behave the same on CPU and GPU.
 
 #### data_format
 

tensorflow_addons/activations/tests/gelu_test.py

@@ -36,6 +36,7 @@ def test_gelu(dtype):
     test_utils.assert_allclose_according_to_type(gelu(x, False), expected_result)
 
 
+@pytest.mark.with_device(["cpu", "gpu"])
 @pytest.mark.parametrize("dtype", [np.float32, np.float64])
 @pytest.mark.parametrize("approximate", [True, False])
 def test_same_as_py_func(dtype, approximate):

tensorflow_addons/conftest.py

@@ -1,9 +1,13 @@
-from tensorflow_addons.utils.test_utils import maybe_run_functions_eagerly  # noqa: F401
-from tensorflow_addons.utils.test_utils import cpu_and_gpu  # noqa: F401
-from tensorflow_addons.utils.test_utils import data_format  # noqa: F401
-from tensorflow_addons.utils.test_utils import set_seeds  # noqa: F401
-from tensorflow_addons.utils.test_utils import pytest_addoption  # noqa: F401
-from tensorflow_addons.utils.test_utils import set_global_variables  # noqa: F401
+from tensorflow_addons.utils.test_utils import (  # noqa: F401
+    maybe_run_functions_eagerly,
+    data_format,
+    set_seeds,
+    pytest_addoption,
+    set_global_variables,
+    pytest_configure,
+    _device_placement,
+    pytest_generate_tests,
+)
 
 # fixtures present in this file will be available
 # when running tests and can be referenced with strings

tensorflow_addons/image/tests/resampler_ops_test.py

@@ -92,7 +92,7 @@ def _make_warp(batch_size, warp_height, warp_width, dtype):
     return warp.astype(dtype)
 
 
-@pytest.mark.usefixtures("cpu_and_gpu")
+@pytest.mark.with_device(["cpu", "gpu"])
 @pytest.mark.parametrize("dtype", [np.float16, np.float32, np.float64])
 def test_op_forward_pass(dtype):
     np.random.seed(0)
@@ -182,7 +182,7 @@ def test_op_errors():
         resampler_ops.resampler(data, warp)
 
 
-@pytest.mark.usefixtures("cpu_and_gpu")
+@pytest.mark.with_device(["cpu", "gpu"])
 @pytest.mark.parametrize("dtype", [np.float16, np.float32, np.float64])
 def test_op_backward_pass(dtype):
     np.random.seed(13)

tensorflow_addons/image/tests/transform_ops_test.py

@@ -19,7 +19,6 @@
 import tensorflow as tf
 
 from tensorflow_addons.image import transform_ops
-from tensorflow_addons.utils import test_utils
 from skimage import transform
 
 _DTYPES = {
@@ -32,39 +31,39 @@
 }
 
 
+@pytest.mark.with_device(["cpu", "gpu"])
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
 @pytest.mark.parametrize("dtype", _DTYPES)
 def test_compose(dtype):
-    with test_utils.use_gpu():
-        image = tf.constant(
-            [[1, 1, 1, 0], [1, 0, 0, 0], [1, 1, 1, 0], [0, 0, 0, 0]], dtype=dtype,
-        )
-        # Rotate counter-clockwise by pi / 2.
-        rotation = transform_ops.angles_to_projective_transforms(np.pi / 2, 4, 4)
-        # Translate right by 1 (the transformation matrix is always inverted,
-        # hence the -1).
-        translation = tf.constant([1, 0, -1, 0, 1, 0, 0, 0], dtype=tf.dtypes.float32)
-        composed = transform_ops.compose_transforms([rotation, translation])
-        image_transformed = transform_ops.transform(image, composed)
-        np.testing.assert_equal(
-            [[0, 0, 0, 0], [0, 1, 0, 1], [0, 1, 0, 1], [0, 1, 1, 1]],
-            image_transformed.numpy(),
-        )
+    image = tf.constant(
+        [[1, 1, 1, 0], [1, 0, 0, 0], [1, 1, 1, 0], [0, 0, 0, 0]], dtype=dtype,
+    )
+    # Rotate counter-clockwise by pi / 2.
+    rotation = transform_ops.angles_to_projective_transforms(np.pi / 2, 4, 4)
+    # Translate right by 1 (the transformation matrix is always inverted,
+    # hence the -1).
+    translation = tf.constant([1, 0, -1, 0, 1, 0, 0, 0], dtype=tf.dtypes.float32)
+    composed = transform_ops.compose_transforms([rotation, translation])
+    image_transformed = transform_ops.transform(image, composed)
+    np.testing.assert_equal(
+        [[0, 0, 0, 0], [0, 1, 0, 1], [0, 1, 0, 1], [0, 1, 1, 1]],
+        image_transformed.numpy(),
+    )
 
 
+@pytest.mark.with_device(["cpu", "gpu"])
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
 @pytest.mark.parametrize("dtype", _DTYPES)
 def test_extreme_projective_transform(dtype):
-    with test_utils.use_gpu():
-        image = tf.constant(
-            [[1, 0, 1, 0], [0, 1, 0, 1], [1, 0, 1, 0], [0, 1, 0, 1]], dtype=dtype,
-        )
-        transformation = tf.constant([1, 0, 0, 0, 1, 0, -1, 0], tf.dtypes.float32)
-        image_transformed = transform_ops.transform(image, transformation)
-        np.testing.assert_equal(
-            [[1, 0, 0, 0], [0, 0, 0, 0], [1, 0, 0, 0], [0, 0, 0, 0]],
-            image_transformed.numpy(),
-        )
+    image = tf.constant(
+        [[1, 0, 1, 0], [0, 1, 0, 1], [1, 0, 1, 0], [0, 1, 0, 1]], dtype=dtype,
+    )
+    transformation = tf.constant([1, 0, 0, 0, 1, 0, -1, 0], tf.dtypes.float32)
+    image_transformed = transform_ops.transform(image, transformation)
+    np.testing.assert_equal(
+        [[1, 0, 0, 0], [0, 0, 0, 0], [1, 0, 0, 0], [0, 0, 0, 0]],
+        image_transformed.numpy(),
+    )
 
 
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")

tensorflow_addons/layers/tests/optical_flow_test.py

@@ -72,7 +72,8 @@ def _create_test_data(data_format):
     return val_a, val_b
 
 
-@pytest.mark.usefixtures("cpu_and_gpu", "maybe_run_functions_eagerly")
+@pytest.mark.with_device(["cpu", "gpu"])
+@pytest.mark.usefixtures("maybe_run_functions_eagerly")
 def test_forward_simple(data_format):
     # We are just testing where the output has vanishing values.
     val_a, val_b = _create_test_data(data_format)
@@ -114,7 +115,7 @@ def test_forward_simple(data_format):
     assert actual.shape == (2, 9, 7, 8)
 
 
-@pytest.mark.usefixtures("cpu_and_gpu")
+@pytest.mark.with_device(["cpu", "gpu"])
 def test_gradients(data_format):
     batch, channels, height, width = 2, 3, 5, 6
     input_a = np.random.randn(batch, channels, height, width).astype(np.float32)
@@ -150,7 +151,7 @@ def correlation_fn(input_a, input_b):
     np.testing.assert_allclose(theoretical[0], numerical[0], atol=1e-3)
 
 
-@pytest.mark.usefixtures("cpu_and_gpu")
+@pytest.mark.with_device(["cpu", "gpu"])
 def test_keras(data_format):
     # Unable to use `layer_test` as this layer has multiple inputs.
     val_a, val_b = _create_test_data(data_format)

tensorflow_addons/metrics/tests/metrics_test.py

@@ -12,18 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-import unittest
 import inspect
 
 from tensorflow.keras.metrics import Metric
 from tensorflow_addons import metrics
 
 
-class MetricsTests(unittest.TestCase):
-    def test_update_state_signature(self):
-        for name, obj in inspect.getmembers(metrics):
-            if inspect.isclass(obj) and issubclass(obj, Metric):
-                check_update_state_signature(obj)
+def test_update_state_signature():
+    for name, obj in inspect.getmembers(metrics):
+        if inspect.isclass(obj) and issubclass(obj, Metric):
+            check_update_state_signature(obj)
 
 
 def check_update_state_signature(metric_class):

tensorflow_addons/optimizers/tests/lazy_adam_test.py

@@ -97,7 +97,7 @@ def _test_sparse(dtype):
 
 
 @pytest.mark.parametrize("dtype", [tf.int32, tf.int64])
-@pytest.mark.usefixtures("cpu_and_gpu")
+@pytest.mark.with_device(["cpu", "gpu"])
 def test_sparse_device_placement(dtype):
 
     # If a GPU is available, tests that all optimizer ops can be placed on

tensorflow_addons/seq2seq/tests/beam_search_ops_test.py

@@ -56,21 +56,18 @@ def test_bad_parent_values_on_cpu():
         [[[0, 0, 0], [0, -1, 1], [2, 1, 2], [-1, -1, -1]]]
     )
     max_sequence_lengths = [3]
-    with tf.device("/cpu:0"):
-        with pytest.raises(tf.errors.InvalidArgumentError):
-            _ = gather_tree(
-                step_ids=step_ids,
-                parent_ids=parent_ids,
-                max_sequence_lengths=max_sequence_lengths,
-                end_token=end_token,
-            )
+
+    with pytest.raises(tf.errors.InvalidArgumentError):
+        _ = gather_tree(
+            step_ids=step_ids,
+            parent_ids=parent_ids,
+            max_sequence_lengths=max_sequence_lengths,
+            end_token=end_token,
+        )
 
 
+@pytest.mark.with_device(["gpu"])
 def test_bad_parent_values_on_gpu():
-    # Only want to run this test on CUDA devices, as gather_tree is not
-    # registered for SYCL devices.
-    if not tf.test.is_gpu_available(cuda_only=True):
-        return
     # (max_time = 4, batch_size = 1, beams = 3)
     # bad parent in beam 1 time 1; appears as a negative index at time 0
     end_token = 10
@@ -82,14 +79,13 @@ def test_bad_parent_values_on_gpu():
     expected_result = _transpose_batch_time(
         [[[2, -1, 2], [6, 5, 6], [7, 8, 9], [10, 10, 10]]]
     )
-    with tf.device("/device:GPU:0"):
-        beams = gather_tree(
-            step_ids=step_ids,
-            parent_ids=parent_ids,
-            max_sequence_lengths=max_sequence_lengths,
-            end_token=end_token,
-        )
-        np.testing.assert_equal(expected_result, beams.numpy())
+    beams = gather_tree(
+        step_ids=step_ids,
+        parent_ids=parent_ids,
+        max_sequence_lengths=max_sequence_lengths,
+        end_token=end_token,
+    )
+    np.testing.assert_equal(expected_result, beams.numpy())
 
 
 def test_gather_tree_batch():