CopyToΒΆ

class dgl.graphbolt.CopyTo(datapipe, device, extra_attrs=None)[source]ΒΆ

Bases: torch.utils.data.dataset.IterableDataset[torch.utils.data.datapipes.datapipe.T_co]

DataPipe that transfers each element yielded from the previous DataPipe to the given device. For MiniBatch, only the related attributes (automatically inferred) will be transferred by default. If you want to transfer any other attributes, indicate them in the extra_attrs.

Functional name: copy_to.

When data has to method implemented, CopyTo will be equivalent to

for data in datapipe:
    yield data.to(device)

For MiniBatch, only a part of attributes will be transferred to accelerate the process by default:

  • When seed_nodes is not None and node_pairs is None, node related

task is inferred. Only labels, sampled_subgraphs, node_features and edge_features will be transferred.

  • When node_pairs is not None and seed_nodes is None, edge/link

related task is inferred. Only labels, compacted_node_pairs, compacted_negative_srcs, compacted_negative_dsts, sampled_subgraphs, node_features and edge_features will be transferred.

  • Otherwise, all attributes will be transferred.

  • If you want some other attributes to be transferred as well, please

specify the name in the extra_attrs. For instance, the following code will copy seed_nodes to the GPU as well:

datapipe = datapipe.copy_to(device="cuda", extra_attrs=["seed_nodes"])
Parameters

  • datapipe (DataPipe) – The DataPipe.

  • device (torch.device) – The PyTorch CUDA device.

  • extra_attrs (List[string]) – The extra attributes of the data in the DataPipe you want to be carried to the specific device. The attributes specified in the extra_attrs will be transferred regardless of the task inferred. It could also be applied to classes other than MiniBatch.