转载请注明出处:小锋学长生活大爆炸[xfxuezhang.cn]
实际上官方的函数解释中就已经非常详细了。
函数形式
代码语言:javascript复制def 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')
函数作用
为分布式训练对图形进行分区,并将分区存储在文件中。
函数内容
分区分为三个步骤: 1) 运行分区算法(如 Metis)将节点分配到分区中; 2) 根据节点分配构建分区图结构; 3) 根据分区结果分割节点特征和边特征。
在对图进行分区时,每个分区都可能包含HALO节点,这些节点被分配给其他分区,但为了提高效率而被包含在本分区中。 在本文中,local nodes/edges指的是真正属于某个分区的节点和边。其余的都是HALO nodes/edges。
分区数据存储在多个文件中,组织结构如下:
代码语言:javascript复制data_root_dir/
|-- graph_name.json # partition configuration file in JSON
|-- node_map.npy # partition id of each node stored in a numpy array (optional)
|-- edge_map.npy # partition id of each edge stored in a numpy array (optional)
|-- part0/ # data for partition 0
|-- node_feats.dgl # node features stored in binary format
|-- edge_feats.dgl # edge features stored in binary format
|-- graph.dgl # graph structure of this partition stored in binary format
|-- part1/ # data for partition 1
|-- node_feats.dgl
|-- edge_feats.dgl
|-- graph.dgl
首先,原始图形和分区的元数据存储在一个以"graph_name"命名的 JSON 文件中。 该 JSON 文件包含原始图的信息以及存储每个分区的文件路径。 下面是一个示例。
代码语言:javascript复制{
"graph_name" : "test",
"part_method" : "metis",
"num_parts" : 2,
"halo_hops" : 1,
"node_map": {
"_U": [ [ 0, 1261310 ],
[ 1261310, 2449029 ] ]
},
"edge_map": {
"_V": [ [ 0, 62539528 ],
[ 62539528, 123718280 ] ]
},
"etypes": { "_V": 0 },
"ntypes": { "_U": 0 },
"num_nodes" : 1000000,
"num_edges" : 52000000,
"part-0" : {
"node_feats" : "data_root_dir/part0/node_feats.dgl",
"edge_feats" : "data_root_dir/part0/edge_feats.dgl",
"part_graph" : "data_root_dir/part0/graph.dgl",
},
"part-1" : {
"node_feats" : "data_root_dir/part1/node_feats.dgl",
"edge_feats" : "data_root_dir/part1/edge_feats.dgl",
"part_graph" : "data_root_dir/part1/graph.dgl",
},
}
- graph_name:是用户给出的图形名称。
- part_method:是将节点分配到分区的方法。目前,支持 "random "和 "metis"。
- num_parts:是分区的数量。
- halo_hops:是分区中作为 HALO 节点的节点跳数。
- node_map:是节点分配映射表,它显示了节点被分配到的分区 ID。
- edge_map:是边的分配映射,它告诉我们边被分配到的分区 ID。
- num_nodes:是全局图中的节点数。
- num_edges:是全局图中的边数。
- part-*:存储一个分区的数据。
如果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``中的信息还可用于计算节点类型和边类型。 "node_map"和"edge_map"中的数据格式如下:
代码语言:javascript复制{
"node_type": [ [ part1_start, part1_end ],
[ part2_start, part2_end ],
... ],
...
},
本质上,``node_map``和`edge_map``是字典。键是节点/边类型。值是成对的列表,包含分区中相应类型的 ID 范围的起点和终点。列表的长度是分区的数量;列表中的每个元素都是一个元组,存储了分区中特定节点/边类型的 ID 范围的起点和终点。
分区的图结构以DGLGraph格式存储在文件中。每个分区中的节点都经过*relabeled*,始终以0开头。我们将原始图中的节点ID称为 "global ID",而将每个分区中重新标记的 ID 称为 "local ID"。每个分区图都有一个整数节点数据张量,存储名为 "dgl.NID",每个值都是节点的全局 ID。同样,边也被重新标记,本地 ID 到全局 ID 的映射被存储为名为 `dgl.EID` 的整数边数据张量。对于异构图,DGLGraph 还包含表示节点类型的节点数据 `dgl.NTYPE`和表示边类型的边数据`dgl.ETYPE`。
分区图包含额外的节点数据("inner_node "和 "orig_id")和边数据("inner_edge"):
- inner_node:表示节点是否属于某个分区。
- inner_edge:表示一条边是否属于一个分区。
- orig_id:在 reshuffle=True 时存在。它表示重新洗牌前原始图中的原始节点 ID。
节点和边的特征被分割开来,与每个图形分区一起存储。分区中的所有节点/边特征都以 DGL 格式存储在一个文件中。节点/边特征存储在字典中,其中键是节点/边数据的名称,值是张量。我们不存储 HALO 节点和边的特征。
在执行 Metis 分区时,我们可以对分区施加一些约束。目前,它支持两种平衡分区的约束条件。默认情况下,Metis 总是尝试平衡每个分区中的节点数。
- balance_ntypes:平衡每个分区中不同类型节点的数量。
- balance_edges:平衡每个分区中的边数。
为了平衡节点类型,用户需要传递一个包含 N 个元素的向量来表示每个节点的类型。N 是输入图中的节点数。
函数入参
- g : DGLGraph
要分割的输入图
- graph_name : str
图的名称。该名称将用于构建 dgl.distributed.DistGraph
- num_parts : int
分区数
- out_path : str
存储所有分区数据文件的路径
- num_hops : int, optional
我们在分区图结构上构建的 HALO 节点的跳数。默认值为 1
- part_method : str, optional
分区方法。支持 "random"和 "metis"。默认值为 "metis"
- reshuffle : bool, optional
是否打乱节点和边,使分区中的节点和边处于连续的 ID 范围内。默认值为 True。该参数已被弃用,将在下一版本中删除
- balance_ntypes : tensor, optional
每个节点的节点类型。这是一个一维整数数组。其值表示每个节点的节点类型。Metis分区使用此参数。指定该参数后,Metis 算法将尝试把输入图分割成多个分区,每个分区中每个节点类型的节点数大致相同。默认值为 "None",这意味着 Metis 只对图进行分区,以平衡节点数量。
- balance_edges : bool
指示是否平衡每个分区中的边。该参数用于 Metis 算法使用。
- return_mapping : bool
如果 `reshuffle=True` 表示返回洗牌后的节点/边 ID 与原始节点/边 ID 之间的映射。
- num_trainers_per_machine : int, optional
每台机器的trainer数量。如果不是 1,则会先将整个图划分给每个trainer,即 num_parts*num_trainers_per_machine parts。每个节点的trainer ID 将存储在节点特征 "trainer_id "中。然后,同一台机器上trainer的分区将被合并成一个更大的分区。分区的最终数量为 "num_part"。
- objtype : str, "cut" or "vol"
将目标设置为边缘切割最小化或通信量最小化。Metis 算法会使用这一参数。
函数返参
- Tensor or dict of tensors, optional
如果 "return_mapping=True",则返回一个一维张量,表示同构图中经过洗牌的节点 ID 与原始节点 ID 之间的映射;如果是异构图,则返回一个一维张量的 dict,其 key 是节点类型,value 是每个节点类型的经过洗牌的节点 ID 与原始节点 ID 之间的一维张量映射。
- Tensor or dict of tensors, optional
如果 "return_mapping=True",则返回一个一维张量,表示同质图中经过洗牌的边 ID 与原始边 ID 之间的映射;如果是异质图,则返回一个一维张量的 dict,其 key 是边类型,value 是每个边类型的经过洗牌的边 ID 与原始边 ID 之间的 1D 张量映射。
使用示例
代码语言:javascript复制>>> dgl.distributed.partition_graph(g, 'test', 4, num_hops=1, part_method='metis',
out_path='output/', reshuffle=True,
balance_ntypes=g.ndata['train_mask'],
balance_edges=True)
>>> g, node_feats, edge_feats, gpb, graph_name = dgl.distributed.load_partition(
'output/test.json', 0)