[GraphBolt] Datapipe for copying to given device (#5918)

python/dgl/graphbolt/__init__.py

@@ -9,6 +9,7 @@ from .graph_storage import *
 from .itemset import *
 from .minibatch_sampler import *
 from .feature_store import *
+from .copy_to import *
 from .dataset import *
 from .subgraph_sampler import *
 

python/dgl/graphbolt/copy_to.py

+"""Graph Bolt CUDA-related Data Pipelines"""
+
+from torchdata.datapipes.iter import IterDataPipe
+
+from ..utils import recursive_apply
+
+
+def _to(x, device):
+    return x.to(device) if hasattr(x, "to") else x
+
+
+class CopyTo(IterDataPipe):
+    """DataPipe that transfers each element yielded from the previous DataPipe
+    to the given device.
+
+    This is equivalent to
+
+    .. code:: python
+
+       for data in datapipe:
+           yield data.to(device)
+
+    Parameters
+    ----------
+    datapipe : DataPipe
+        The DataPipe.
+    device : torch.device
+        The PyTorch CUDA device.
+    """
+
+    def __init__(self, datapipe, device):
+        super().__init__()
+        self.datapipe = datapipe
+        self.device = device
+
+    def __iter__(self):
+        for data in self.datapipe:
+            data = recursive_apply(data, _to, self.device)
+            yield data

tests/python/pytorch/graphbolt/test_copy_to.py

+import unittest
+
+import backend as F
+import dgl.graphbolt
+import torch
+
+
+@unittest.skipIf(F._default_context_str == "cpu", "CopyTo needs GPU to test")
+def test_CopyTo():
+    dp = dgl.graphbolt.MinibatchSampler(torch.randn(20), 4)
+    dp = dgl.graphbolt.CopyTo(dp, "cuda")
+
+    for data in dp:
+        assert data.device.type == "cuda"