dgl.DGLGraph.pin_memory_

DGLGraph.pin_memory_()

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

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

对于已固定的图，不允许实例化新的稀疏格式。为了避免在训练期间隐式实例化格式，你应该在固定之前创建所有需要的格式。但克隆和实例化是可以的。参见下面的示例。

返回:

已固定的图。

返回类型:

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])