dgl.DGLGraph.to

DGLGraph.to(device, **kwargs)

将 ndata、edata 和图结构移动到目标设备 (CPU/GPU)。

如果图已在指定的设备上，该函数将直接返回它。否则，它将返回在指定设备上的图的克隆。

请注意，节点和边特征的数据在被访问或调用 materialize_data() 之前不会移动到指定设备上。

参数:

• device (框架特定的设备上下文对象) – 数据要移动到的上下文（例如，torch.device）。

• kwargs (关键字参数。) – 传递给框架复制函数的关键字参数。

返回:

在指定设备上的图。

返回类型:

DGLGraph

示例

以下示例使用 PyTorch 后端。

>>> import dgl
>>> import torch

>>> g = dgl.graph((torch.tensor([1, 0]), torch.tensor([1, 2])))
>>> g.ndata['h'] = torch.ones(3, 1)
>>> g.edata['h'] = torch.zeros(2, 2)
>>> g1 = g.to(torch.device('cuda:0'))
>>> print(g1.device)
device(type='cuda', index=0)
>>> print(g1.ndata['h'].device)
device(type='cuda', index=0)
>>> print(g1.nodes().device)
device(type='cuda', index=0)

原始图仍在 CPU 上。

>>> print(g.device)
device(type='cpu')
>>> print(g.ndata['h'].device)
device(type='cpu')
>>> print(g.nodes().device)
device(type='cpu')

异构图的情况相同。