Refactor TFDataLayer to be more generic for Grain. (#21598)

Author: james77777778

keras/src/backend/torch/core.py

@@ -191,21 +191,18 @@ def convert_to_tensor(x, dtype=None, sparse=None, ragged=None):
         raise ValueError("`sparse=True` is not supported with torch backend")
     if ragged:
         raise ValueError("`ragged=True` is not supported with torch backend")
-    if isinstance(x, Variable):
-        if dtype is None:
-            return x.value
-        x = x.value
-        return x.to(to_torch_dtype(dtype))
-    if is_tensor(x):
+    if isinstance(x, Variable) or is_tensor(x):
+        if isinstance(x, Variable):
+            x = x.value
         device = get_device()
         if x.device != device:
             if x.is_meta:
                 x = torch.empty_like(x, device=device)
             else:
                 x = x.to(device)
-        if dtype is None:
-            return x
-        return x.to(to_torch_dtype(dtype))
+        if dtype is not None:
+            x = x.to(to_torch_dtype(dtype))
+        return x
     if dtype is None:
         if isinstance(x, bool):
             return torch.as_tensor(x, dtype=torch.bool, device=get_device())

keras/src/layers/preprocessing/category_encoding.py

@@ -1,12 +1,12 @@
 from keras.src.api_export import keras_export
 from keras.src.backend import KerasTensor
-from keras.src.layers.preprocessing.tf_data_layer import TFDataLayer
+from keras.src.layers.preprocessing.data_layer import DataLayer
 from keras.src.utils import backend_utils
 from keras.src.utils import numerical_utils
 
 
 @keras_export("keras.layers.CategoryEncoding")
-class CategoryEncoding(TFDataLayer):
+class CategoryEncoding(DataLayer):
     """A preprocessing layer which encodes integer features.
 
     This layer provides options for condensing data into a categorical encoding
@@ -15,7 +15,7 @@ class CategoryEncoding(TFDataLayer):
     inputs. For integer inputs where the total number of tokens is not known,
     use `keras.layers.IntegerLookup` instead.
 
-    **Note:** This layer is safe to use inside a `tf.data` pipeline
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
     (independently of which backend you're using).
 
     Examples:

keras/src/layers/preprocessing/data_layer.py

@@ -0,0 +1,159 @@
+import keras.src.backend
+from keras.src import tree
+from keras.src.layers.layer import Layer
+from keras.src.random.seed_generator import SeedGenerator
+from keras.src.utils import backend_utils
+from keras.src.utils import jax_utils
+from keras.src.utils import tracking
+
+
+class DataLayer(Layer):
+    """Layer designed for safe use in `tf.data` or `grain` pipeline.
+
+    This layer overrides the `__call__` method to ensure that the correct
+    backend is used and that computation is performed on the CPU.
+
+    The `call()` method in subclasses should use `self.backend` ops. If
+    randomness is needed, define both `seed` and `generator` in `__init__` and
+    retrieve the running seed using `self._get_seed_generator()`. If the layer
+    has weights in `__init__` or `build()`, use `convert_weight()` to ensure
+    they are in the correct backend.
+
+    **Note:** This layer and its subclasses only support a single input tensor.
+
+    Examples:
+
+    **Custom `DataLayer` subclass:**
+
+    ```python
+    from keras.src.layers.preprocessing.data_layer import DataLayer
+    from keras.src.random import SeedGenerator
+
+
+    class BiasedRandomRGBToHSVLayer(DataLayer):
+        def __init__(self, seed=None, **kwargs):
+            super().__init__(**kwargs)
+            self.probability_bias = ops.convert_to_tensor(0.01)
+            self.seed = seed
+            self.generator = SeedGenerator(seed)
+
+        def call(self, inputs):
+            images_shape = self.backend.shape(inputs)
+            batch_size = 1 if len(images_shape) == 3 else images_shape[0]
+            seed = self._get_seed_generator(self.backend._backend)
+
+            probability = self.backend.random.uniform(
+                shape=(batch_size,),
+                minval=0.0,
+                maxval=1.0,
+                seed=seed,
+            )
+            probability = self.backend.numpy.add(
+                probability, self.convert_weight(self.probability_bias)
+            )
+            hsv_images = self.backend.image.rgb_to_hsv(inputs)
+            return self.backend.numpy.where(
+                probability[:, None, None, None] > 0.5,
+                hsv_images,
+                inputs,
+            )
+
+        def compute_output_shape(self, input_shape):
+            return input_shape
+    ```
+
+    **Using as a regular Keras layer:**
+
+    ```python
+    import numpy as np
+
+    x = np.random.uniform(size=(1, 16, 16, 3)).astype("float32")
+    print(BiasedRandomRGBToHSVLayer()(x).shape)  # (1, 16, 16, 3)
+    ```
+
+    **Using in a `tf.data` pipeline:**
+
+    ```python
+    import tensorflow as tf
+
+    tf_ds = tf.data.Dataset.from_tensors(x)
+    tf_ds = tf_ds.map(BiasedRandomRGBToHSVLayer())
+    print([x.shape for x in tf_ds])  # [(1, 16, 16, 3)]
+    ```
+
+    **Using in a `grain` pipeline:**
+
+    ```python
+    import grain
+
+    grain_ds = grain.MapDataset.source([x])
+    grain_ds = grain_ds.map(BiasedRandomRGBToHSVLayer())
+    print([x.shape for x in grain_ds])  # [(1, 16, 16, 3)]
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.backend = backend_utils.DynamicBackend()
+        self._allow_non_tensor_positional_args = True
+
+    def __call__(self, inputs, **kwargs):
+        sample_input = tree.flatten(inputs)[0]
+        if (
+            not isinstance(sample_input, keras.KerasTensor)
+            and backend_utils.in_tf_graph()
+            and not jax_utils.is_in_jax_tracing_scope(sample_input)
+        ):
+            # We're in a TF graph, e.g. a tf.data pipeline.
+            self.backend.set_backend("tensorflow")
+            inputs = tree.map_structure(
+                lambda x: self.backend.convert_to_tensor(
+                    x, dtype=self.compute_dtype
+                ),
+                inputs,
+            )
+            switch_convert_input_args = False
+            if self._convert_input_args:
+                self._convert_input_args = False
+                switch_convert_input_args = True
+            try:
+                outputs = super().__call__(inputs, **kwargs)
+            finally:
+                self.backend.reset()
+                if switch_convert_input_args:
+                    self._convert_input_args = True
+            return outputs
+        elif (
+            not isinstance(sample_input, keras.KerasTensor)
+            and backend_utils.in_grain_data_pipeline()
+        ):
+            # We're in a Grain data pipeline. Force computation and data
+            # placement to CPU.
+            with keras.src.backend.device_scope("cpu"):
+                return super().__call__(inputs, **kwargs)
+        else:
+            return super().__call__(inputs, **kwargs)
+
+    @tracking.no_automatic_dependency_tracking
+    def _get_seed_generator(self, backend=None):
+        if not hasattr(self, "seed") or not hasattr(self, "generator"):
+            raise ValueError(
+                "The `seed` and `generator` variable must be set in the "
+                "`__init__` method before calling `_get_seed_generator()`."
+            )
+        if backend is None or backend == keras.backend.backend():
+            return self.generator
+        if not hasattr(self, "_backend_generators"):
+            self._backend_generators = {}
+        if backend in self._backend_generators:
+            return self._backend_generators[backend]
+        seed_generator = SeedGenerator(self.seed, backend=self.backend)
+        self._backend_generators[backend] = seed_generator
+        return seed_generator
+
+    def convert_weight(self, weight):
+        """Convert the weight if it is from the a different backend."""
+        if self.backend.name == keras.backend.backend():
+            return weight
+        else:
+            weight = keras.ops.convert_to_numpy(weight)
+            return self.backend.convert_to_tensor(weight)

keras/src/layers/preprocessing/data_layer_test.py

@@ -0,0 +1,90 @@
+import grain
+import numpy as np
+import pytest
+from tensorflow import data as tf_data
+
+from keras.src import backend
+from keras.src import testing
+from keras.src.layers.preprocessing.data_layer import DataLayer
+from keras.src.random import SeedGenerator
+
+
+class RandomRGBToHSVLayer(DataLayer):
+    def __init__(self, data_format=None, seed=None, **kwargs):
+        super().__init__(**kwargs)
+        self.data_format = backend.standardize_data_format(data_format)
+        self.seed = seed
+        self.generator = SeedGenerator(seed)
+
+    def call(self, inputs):
+        images_shape = self.backend.shape(inputs)
+        batch_size = 1 if len(images_shape) == 3 else images_shape[0]
+        seed = self._get_seed_generator(self.backend._backend)
+
+        probability = self.backend.random.uniform(
+            shape=(batch_size,),
+            minval=0.0,
+            maxval=1.0,
+            seed=seed,
+        )
+        hsv_images = self.backend.image.rgb_to_hsv(
+            inputs, data_format=self.data_format
+        )
+        return self.backend.numpy.where(
+            probability[:, None, None, None] > 0.5, hsv_images, inputs
+        )
+
+    def compute_output_shape(self, input_shape):
+        return input_shape
+
+
+class DataLayerTest(testing.TestCase):
+    @pytest.mark.requires_trainable_backend
+    def test_layer(self):
+        self.run_layer_test(
+            RandomRGBToHSVLayer,
+            init_kwargs={
+                "seed": 1337,
+                "data_format": "channels_last",
+            },
+            input_shape=(1, 2, 2, 3),
+            supports_masking=False,
+            expected_output_shape=(1, 2, 2, 3),
+        )
+
+        self.run_layer_test(
+            RandomRGBToHSVLayer,
+            init_kwargs={
+                "seed": 1337,
+                "data_format": "channels_first",
+            },
+            input_shape=(1, 3, 2, 2),
+            supports_masking=False,
+            expected_output_shape=(1, 3, 2, 2),
+        )
+
+    def test_tf_data_compatibility(self):
+        data_format = backend.config.image_data_format()
+        if data_format == "channels_last":
+            input_data = np.random.random((2, 8, 8, 3)).astype("float32")
+        else:
+            input_data = np.random.random((2, 3, 8, 8)).astype("float32")
+        layer = RandomRGBToHSVLayer(data_format=data_format, seed=1337)
+
+        ds = tf_data.Dataset.from_tensor_slices(input_data).batch(2).map(layer)
+        for output in ds.take(1):
+            self.assertDType(output, "float32")
+            self.assertEqual(list(output.shape), list(input_data.shape))
+
+    def test_grain_compatibility(self):
+        data_format = backend.config.image_data_format()
+        if data_format == "channels_last":
+            input_data = np.random.random((2, 8, 8, 3)).astype("float32")
+        else:
+            input_data = np.random.random((2, 3, 8, 8)).astype("float32")
+        layer = RandomRGBToHSVLayer(data_format=data_format, seed=1337)
+
+        ds = grain.MapDataset.source(input_data).batch(2).map(layer)
+        for output in ds[:1]:
+            self.assertDType(output, "float32")
+            self.assertEqual(list(output.shape), list(input_data.shape))

keras/src/layers/preprocessing/discretization.py

@@ -2,21 +2,21 @@
 
 from keras.src import backend
 from keras.src.api_export import keras_export
-from keras.src.layers.preprocessing.tf_data_layer import TFDataLayer
+from keras.src.layers.preprocessing.data_layer import DataLayer
 from keras.src.utils import argument_validation
 from keras.src.utils import numerical_utils
 from keras.src.utils.module_utils import tensorflow as tf
 
 
 @keras_export("keras.layers.Discretization")
-class Discretization(TFDataLayer):
+class Discretization(DataLayer):
     """A preprocessing layer which buckets continuous features by ranges.
 
     This layer will place each element of its input data into one of several
     contiguous ranges and output an integer index indicating which range each
     element was placed in.
 
-    **Note:** This layer is safe to use inside a `tf.data` pipeline
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
     (independently of which backend you're using).
 
     Input shape:

keras/src/layers/preprocessing/feature_space.py

@@ -3,7 +3,7 @@
 from keras.src import tree
 from keras.src.api_export import keras_export
 from keras.src.layers.layer import Layer
-from keras.src.layers.preprocessing.tf_data_layer import TFDataLayer
+from keras.src.layers.preprocessing.data_layer import DataLayer
 from keras.src.saving import saving_lib
 from keras.src.saving import serialization_lib
 from keras.src.saving.keras_saveable import KerasSaveable
@@ -723,7 +723,7 @@ def __call__(self, data):
                 data[name] = tf.expand_dims(x, -1)
 
         with backend_utils.TFGraphScope():
-            # This scope is to make sure that inner TFDataLayers
+            # This scope is to make sure that inner DataLayers
             # will not convert outputs back to backend-native --
             # they should be TF tensors throughout
             preprocessed_data = self._preprocess_features(data)
@@ -808,7 +808,7 @@ def load_own_variables(self, store):
         return
 
 
-class TFDConcat(TFDataLayer):
+class TFDConcat(DataLayer):
     def __init__(self, axis, **kwargs):
         super().__init__(**kwargs)
         self.axis = axis
@@ -817,6 +817,6 @@ def call(self, xs):
         return self.backend.numpy.concatenate(xs, axis=self.axis)
 
 
-class TFDIdentity(TFDataLayer):
+class TFDIdentity(DataLayer):
     def call(self, x):
         return x

keras/src/layers/preprocessing/image_preprocessing/aug_mix.py

@@ -43,6 +43,13 @@ class AugMix(BaseImagePreprocessingLayer):
     in num_chains different ways, with each chain consisting of
     chain_depth augmentations.
 
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
+    (independently of which backend you're using).
+
+    References:
+        - [AugMix paper](https://arxiv.org/pdf/1912.02781)
+        - [Official Code](https://github.com/google-research/augmix)
+
     Args:
         value_range: the range of values the incoming images will have.
             Represented as a two number tuple written (low, high).
@@ -64,10 +71,6 @@ class AugMix(BaseImagePreprocessingLayer):
         interpolation: The interpolation method to use for resizing operations.
             Options include `"nearest"`, `"bilinear"`. Default is `"bilinear"`.
         seed: Integer. Used to create a random seed.
-
-    References:
-        - [AugMix paper](https://arxiv.org/pdf/1912.02781)
-        - [Official Code](https://github.com/google-research/augmix)
     """
 
     _USE_BASE_FACTOR = False

keras/src/layers/preprocessing/image_preprocessing/auto_contrast.py

@@ -17,6 +17,9 @@ class AutoContrast(BaseImagePreprocessingLayer):
 
     This layer is active at both training and inference time.
 
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
+    (independently of which backend you're using).
+
     Args:
         value_range: Range of values the incoming images will have.
             Represented as a two number tuple written `(low, high)`.