Force computation and data on CPU when in grain data pipeline. (#21553)

Author: james77777778

keras/src/backend/jax/distribution_lib.py

@@ -201,7 +201,11 @@ def process_id():
 def _to_backend_device(device_name):
     if isinstance(device_name, jax.Device):
         return device_name
-    device_type, device_id = device_name.split(":")
+    device_name = str(device_name)
+    if ":" not in device_name:
+        device_type, device_id = device_name, 0
+    else:
+        device_type, device_id = device_name.split(":")
 
     devices = jax.devices(backend=device_type)
     for device in devices:

keras/src/backend/torch/core.py

@@ -62,7 +62,7 @@ def device_scope(device_name):
     current_device = _parse_device_input(device_name)
     global_state.set_global_attribute("torch_device", current_device)
     try:
-        yield
+        yield torch.device(current_device)
     finally:
         global_state.set_global_attribute("torch_device", previous_device)
 

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

@@ -1,3 +1,4 @@
+import grain
 import numpy as np
 import pytest
 from absl.testing import parameterized
@@ -158,6 +159,34 @@ def test_tf_data_compatibility(self):
         output = next(iter(ds)).numpy()
         self.assertEqual(tuple(output.shape), output_shape)
 
+    def test_grain_compatibility(self):
+        if backend.config.image_data_format() == "channels_last":
+            input_shape = (2, 10, 12, 3)
+            output_shape = (2, 8, 9, 3)
+        else:
+            input_shape = (2, 3, 10, 12)
+            output_shape = (2, 3, 8, 9)
+        layer = layers.Resizing(8, 9)
+        input_data = np.random.random(input_shape)
+        ds = (
+            grain.MapDataset.source(input_data)
+            .to_iter_dataset()
+            .batch(2)
+            .map(layer)
+        )
+        output = next(iter(ds))
+        output_np = backend.convert_to_numpy(output)
+
+        self.assertEqual(tuple(output_np.shape), output_shape)
+        self.assertTrue(backend.is_tensor(output))
+        # Ensure the device of the data is on CPU.
+        if backend.backend() == "tensorflow":
+            self.assertIn("CPU", str(output.device))
+        elif backend.backend() == "jax":
+            self.assertIn("CPU", str(output.device))
+        elif backend.backend() == "torch":
+            self.assertEqual("cpu", str(output.device))
+
     @pytest.mark.skipif(
         backend.backend() != "tensorflow",
         reason="Sequential + tf.data only works with TF backend",

keras/src/layers/preprocessing/rescaling_test.py

@@ -1,3 +1,4 @@
+import grain
 import numpy as np
 import pytest
 from tensorflow import data as tf_data
@@ -74,6 +75,21 @@ def test_tf_data_compatibility(self):
         ds = tf_data.Dataset.from_tensor_slices(x).batch(3).map(layer)
         next(iter(ds)).numpy()
 
+    def test_grain_compatibility(self):
+        layer = layers.Rescaling(scale=1.0 / 255, offset=0.5)
+        x = np.random.random((3, 10, 10, 3)) * 255
+        ds = grain.MapDataset.source(x).to_iter_dataset().batch(3).map(layer)
+        output = next(iter(ds))
+
+        self.assertTrue(backend.is_tensor(output))
+        # Ensure the device of the data is on CPU.
+        if backend.backend() == "tensorflow":
+            self.assertIn("CPU", str(output.device))
+        elif backend.backend() == "jax":
+            self.assertIn("CPU", str(output.device))
+        elif backend.backend() == "torch":
+            self.assertEqual("cpu", str(output.device))
+
     def test_rescaling_with_channels_first_and_vector_scale(self):
         config = backend.image_data_format()
         backend.set_image_data_format("channels_first")

keras/src/layers/preprocessing/tf_data_layer.py

@@ -46,7 +46,16 @@ def __call__(self, inputs, **kwargs):
                 if switch_convert_input_args:
                     self._convert_input_args = True
             return outputs
-        return super().__call__(inputs, **kwargs)
+        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):

keras/src/utils/backend_utils.py

@@ -1,5 +1,6 @@
 import copy
 import importlib
+import inspect
 import os
 import sys
 
@@ -40,6 +41,26 @@ def __exit__(self, *args, **kwargs):
         )
 
 
+def in_grain_data_pipeline():
+    if "grain" not in sys.modules:
+        # Fast path to check if grain is not imported.
+        return False
+
+    # We use a lightweight version of `inspect.stack` to detect execution within
+    # grain.
+    current_frame = inspect.currentframe()
+    while current_frame:
+        if (
+            os.path.join("grain", "_src", "python", "dataset")
+            in current_frame.f_code.co_filename
+            or os.path.join("grain", "_src", "python", "data_loader")
+            in current_frame.f_code.co_filename
+        ):
+            return True
+        current_frame = current_frame.f_back
+    return False
+
+
 class DynamicBackend:
     """A class that can be used to switch from one backend to another.