dgl.DGLGraph.pin_memory_
- DGLGraph.pin_memory_()[source]
将图结构和节点/边数据固定到页面锁定内存中,以便GPU进行零拷贝访问。
这是一个原地方法。图结构必须在CPU上才能被固定。 如果图结构已经被固定,函数会直接返回它。
不允许为固定图创建新的稀疏格式。为了避免在训练期间隐式创建格式,您应该在固定之前创建所有需要的格式。但克隆和创建是允许的。请参见以下示例。
- Returns:
固定的图表。
- Return type:
示例
以下示例使用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])