DGL获取节点邻居代码

简单记录一下DGL根据节点采样子图的代码。

先创建一个简单的同构图。

import dgl
import torch

# 边 0->1, 0->2, 0->3, 1->3
u, v = torch.tensor([0, 0, 0, 1]), torch.tensor([1, 2, 3, 3])
g = dgl.graph((u, v))
#节点特征
g.ndata['feat'] = torch.randn(4, 64)
#节点标签
g.ndata['label'] = torch.tensor([0,1,1,0])
print(g)

然后采样得到节点3的所有一阶邻居

nids = [3]
frontier = g.sample_neighbors(nids, -1)
us = frontier.edges()[0].numpy()
vs = frontier.edges()[1].numpy()

print(frontier)
print(us)
print(vs)