dgl.DGLGraph.pin_memory_

DGLGraph.pin_memory_()[source]

将图结构和节点/边数据固定到页面锁定内存中,以便GPU进行零拷贝访问。

这是一个原地方法。图结构必须在CPU上才能被固定。 如果图结构已经被固定,函数会直接返回它。

不允许为固定图创建新的稀疏格式。为了避免在训练期间隐式创建格式,您应该在固定之前创建所有需要的格式。但克隆和创建是允许的。请参见以下示例。

Returns:

固定的图表。

Return type:

DGLGraph

示例

以下示例使用PyTorch后端。

>>> import dgl
>>> import torch
>>> g = dgl.graph((torch.tensor([1, 0]), torch.tensor([1, 2])))
>>> g.pin_memory_()

对于固定的图,不允许实现新的稀疏格式。

>>> g.create_formats_()  # This would raise an error! You should do this before pinning.

允许克隆和物化新格式。返回的图被固定。

>>> g1 = g.formats(['csc'])
>>> assert not g1.is_pinned()

固定的图可以从CPU和GPU访问。具体设备取决于query的上下文。例如,find_edges()中的eid是一个查询。当eid在CPU上时,find_edges()在CPU上执行,并且返回的值是CPU张量。

>>> g.unpin_memory_()
>>> g.create_formats_()
>>> g.pin_memory_()
>>> eid = torch.tensor([1])
>>> g.find_edges(eids)
(tensor([0]), tensor([2]))

eid 移动到 GPU 上,find_edges() 将在 GPU 上执行,并且返回的值是 GPU 张量。

>>> eid = eid.to('cuda:0')
>>> g.find_edges(eids)
(tensor([0], device='cuda:0'), tensor([2], device='cuda:0'))

如果你不提供query,方法将默认在CPU上执行。

>>> g.in_degrees()
tensor([0, 1, 1])