【DGL系列】详细分析DGL中dgl.NID和orig_id的区别

转载请注明出处：小锋学长生活大爆炸[xfxuezhagn.cn]

如果本文帮助到了你，欢迎[点赞、收藏、关注]哦~

目录

背景知识

深入分析

初步结论

代码验证

实验设计

结果分析

最终结论

扩展思考

---

本文将详细分析orig_id和dgl.NID的区别。

背景知识

        在做子图分区的时候，可以返回NID和orig_id，具体我们看看官方教程里的介绍：

以下来自：7.1 Preprocessing for Distributed Training — DGL 0.8.2post1 documentation

---

By default, the partition API assigns new IDs to the nodes and edges in the input graph to help locate nodes/edges during distributed training/inference. After assigning IDs, the partition API shuffles all node data and edge data accordingly. After generating partitioned subgraphs, each subgraph is stored as a DGLGraph object. The original node/edge IDs before reshuffling are stored in the field of ‘orig_id’ in the node/edge data of the subgraphs. The node data dgl.NID and the edge data dgl.EID of the subgraphs store new node/edge IDs of the full graph after nodes/edges reshuffle. During the training, users just use the new node/edge IDs.

1. 默认情况下，分区 API 会为输入图中的节点和边分配新的 ID，以帮助在分布式训练/推理期间定位节点/边。
2. 分配 ID 后，分区 API 会相应地洗牌所有节点数据和边数据。生成分区子图后，每个子图都存储为 DGLGraph 对象。
3. 重新洗牌前的原始节点/边 ID 存储在子图的节点/边数据的“orig_id”字段中。
4. 子图的节点数据 dgl.NID 和边数据 dgl.EID 存储节点/边重新洗牌后完整图的新节点/边 ID。
5. 在训练期间，用户只需使用新的节点/边 ID。

提醒：这里的“重新洗牌 reshuffle”指的是“重新排序”。

深入分析

        上面的大概意思就是说，orig_id存储的是打乱前节点在原本大图的id，而NID存储的是打乱后节点在原本大图的id。

        我们先看一下执行分区的函数partition_graph：

dgl.distributed.partition.partition_graph — DGL 0.8.2post1 documentation

dgl.distributed.partition.partition_graph(g, graph_name, num_parts, out_path, num_hops=1, part_method='metis', reshuffle=True, balance_ntypes=None, balance_edges=False, return_mapping=False, num_trainers_per_machine=1, objtype='cut')

需要注意的是：

        如果 reshuffle=False，则分区的节点 ID 和边 ID 不属于连续的 ID 范围。在这种情况下，DGL 将节点/边映射（从节点/边 ID 到分区 ID）存储在单独的文件（node_map.npy 和 edge_map.npy）中。节点/边映射存储在 numpy 文件中。此格式已弃用，下一个版本将不再支持此格式。换言之，未来版本在对图形进行分区时将始终对节点 ID 和边 ID 进行随机排序。

        如果 reshuffle=True，则 node_map 和 edge_map 包含用于在全局节点/边 ID 到分区本地节点/边 ID 之间映射的信息。对于异构图，node_map和edge_map中的信息也可用于计算节点类型和边类型。该操作可以让分区中的节点和边位于连续的 ID 范围内。

        从本质上讲，node_map 和 edge_map 是字典。键是节点/边缘类型。这些值是包含分区中相应类型的 ID 范围的开始和结束对的列表。列表的长度是分区的数量;列表中的每个元素都是一个元组，用于存储分区中特定节点/边缘类型的 ID 范围的开始和结束。

        分区的图形结构存储在 DGLGraph 格式的文件中。每个分区中的节点都会被重新标记为始终从0开始。我们将原始图中的节点 ID 称为 global ID，而将每个分区中重新标记的 ID 称为 local ID。每个分区图都有一个节点数据张量，存储在名为 dgl.NID 的字段下，其中的每个值都是该节点的全局 ID。同样，边也会被重新标记，从本地 ID 到全局 ID 的映射将存储为名为 dgl.EID 的整数边数据张量下。对于异构图，DGLGraph 还包含一个节点数据 dgl.NTYPE 用于表示节点类型和边数据 dgl.ETYPE 表示边类型。

        当 reshuffle=True 时，“orig_id”存在。它表示reshuffle之前原始图中的原始节点 ID。

初步结论

        上面也就是说，当 reshuffle=True 时，才会返回“orig_id”字段。考虑到分区完，子分区上的节点ID可能是不连续的(可能影响后续算法执行)，所以reshuffle就是重新分配ID，以便在该子分区上的ID能够连续。

        因此，orig_id是洗牌前的大图ID，dgl.NID是洗牌后的大图ID。

代码验证

实验设计

        我们通过简单代码验证下是不是这样，我们以节点N来看。

原本的大图：

# 定义图的边
src_nodes = torch.tensor([0, 1, 2, 3, 4, 2])  # 起始节点
dst_nodes = torch.tensor([1, 2, 3, 4, 5, 4])  # 结束节点

# 创建图对象
g = dgl.graph((src_nodes, dst_nodes))

# 图是无向的，所以添加反向边
g = dgl.to_bidirected(g)

# 打印图的信息
print("Nodes in the graph:", g.nodes())
print("Edges in the graph:", g.edges())
plot_dgl_graph(g)

输出：

Nodes in the graph: tensor([0, 1, 2, 3, 4, 5])
Edges in the graph: (tensor([0, 0, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4, 5]), tensor([1, 3, 0, 2, 1, 3, 4, 0, 2, 4, 2, 3, 5, 4]))

进行分区：

partition = dgl.distributed.partition_graph(g, graph_name='test', num_parts=2,
                                    out_path='./test/', num_hops=1, return_mapping=True,
                                    balance_edges=False) 

获取分区1的信息：

# 读取子图信息
g1, nodes_feats1, efeats1, gpb1, graph_name1, node_type1, etype1 = dgl.distributed.load_partition('./test/test.json', 0)

print(g1.ndata[dgl.NID])
print(g1.ndata['orig_id'])

输出：

nodes:   tensor([0, 1, 2, 3, 4, 5])
NID:      tensor([0, 1, 2, 3, 4, 5])
orig_id: tensor([1, 4, 5, 0, 2, 3])

# 注意，后三个(即nodes中的3,4,5)是halo节点

获取分区2的信息：

# 读取子图信息
g2, nodes_feats2, efeats2, gpb2, graph_name2, node_type2, etype2 = dgl.distributed.load_partition('./test/test.json', 1)

print(g2.ndata[dgl.NID])
print(g2.ndata['orig_id'])

输出：

nodes:   tensor([0, 1, 2, 3, 4])
NID:      tensor([3, 4, 5, 0, 1])
orig_id: tensor([0, 2, 3, 1, 4])

# 注意，后两个(即nodes中的3,4)是halo节点

结果分析

从上面的分区1和分区2的结果上可以看出：

• 每个分区中的g.nodes()都是从0开始的，确实每个分区的节点被重新分配了ID。验证了“背景知识”里的第1、2条；
• 节点并不是按顺序划分到子分区，因此每个分区中的orig_id是不连续的，并且反映了最原始的大图中的节点ID。验证了“背景知识”里的第3条；
• reshuffle操作对大图的节点ID进行了重新排序，因此可以看到每个分区中的NID确实是连续的。验证了“背景知识”里的第4条；

最终结论

        因此，可以有以下结论：

1. orig_id存储的是重新排序前，节点在大图上的ID；
2. dgl.NID存储的是重新排序后，节点在大图上的ID；
3. 两者都是global id；
4. orig_id存储的才是真正的、最原始的节点ID；
5. dgl.NID存储的ID虽然也能代表全局ID，但它是重新排序后的ID；
6. 第4和5点反映出，节点位置如果变化，orig_id不会变，但dgl.NID可能会变；

基于以上几点，在使用的时候需要多加注意区分。正如“背景知识”的第5点所说，我觉得大部分情况下，dgl.NID应该就够用了。

扩展思考

1、你知道gpb1.partid2nids(0)、gpb1.partid2nids(1)返回的是NID还是orig_id吗？

print('partid2nids of part 0: ', gpb1.partid2nids(0))
print('partid2nids of part 1: ', gpb1.partid2nids(1))

输出：

partid2nids of part 0:  tensor([0, 1, 2])
partid2nids of part 1:  tensor([3, 4, 5])

        所以，它返回的是NID哦。

2、通过代码也可以看出来，NID基本上是按照升序排序的，而且内点inner node是在前面，外点halo node是在后面。比如tensor([0, 1, 2, 3, 4, 5])，其中0,1,2就是当前分区的内点，3,4,5就是其他分区在当前分区上的halo点。