Force Input DS on Device

Author: DEKHTIARJonathan

tftrt/examples/benchmark_runner.py

@@ -20,6 +20,7 @@
 from benchmark_utils import timed_section
 
 from dataloading_utils import SyntheticDataset
+from dataloading_utils import ensure_dataset_on_gpu
 from dataloading_utils import get_dequeue_batch_fn
 from dataloading_utils import get_force_data_on_gpu_fn
 
@@ -352,6 +353,8 @@ def execute_benchmark(self):
                         f"synthetic dataset. Performance numbers will be "
                         f"impacted.\nError: {str(e)}."
                     )
+            else:
+                dataset = ensure_dataset_on_gpu(dataset, device="GPU:0")
 
             @force_gpu_resync
             @tf.function(jit_compile=self._args.use_xla)

tftrt/examples/dataloading_utils.py

@@ -36,6 +36,18 @@ def __iter__(self):
             yield data_batch
 
 
+def ensure_dataset_on_gpu(dataset, device):
+    if device.lower() not in dataset._variant_tensor_attr.device.lower():
+        return dataset.apply(
+            tf.data.experimental.prefetch_to_device(
+                device=device,
+                buffer_size=tf.data.experimental.AUTOTUNE
+            )
+        )
+    else:
+        return dataset
+
+
 def get_dequeue_batch_fn(ds_iter):
 
     @force_gpu_resync