dgl.distributed.partition_graph

dgl.distributed.partition_graph(g, graph_name, num_parts, out_path, num_hops=1, part_method='metis', balance_ntypes=None, balance_edges=False, return_mapping=False, num_trainers_per_machine=1, objtype='cut', graph_formats=None, use_graphbolt=False, **kwargs)[source]

为分布式训练分割图并将分割存储在文件中。

分区过程分为三个步骤：1) 运行分区算法（例如，Metis）将节点分配到分区；2) 根据节点分配构建分区图结构；3) 根据分区结果分割节点特征和边特征。

当图被分区时，每个分区可以包含HALO节点，这些节点被分配给其他分区，但为了提高效率而包含在此分区中。在本文档中，本地节点/边指的是真正属于一个分区的节点和边。其余的是“HALO节点/边”。

分区数据被存储到多个文件中，组织方式如下：

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文件包含了原始图的信息以及存储每个分区的文件路径。下面展示了一个示例。

{
   "graph_name" : "test",
   "part_method" : "metis",
   "num_parts" : 2,
   "halo_hops" : 1,
   "node_map": {
       "_N": [ [ 0, 1261310 ],
               [ 1261310, 2449029 ] ]
   },
   "edge_map": {
       "_N:_E:_N": [ [ 0, 62539528 ],
                     [ 62539528, 123718280 ] ]
   },
   "etypes": { "_N:_E:_N": 0 },
   "ntypes": { "_N": 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。 node_map 的格式如下所述。
edge_map 是边分配映射，它告诉边被分配到的分区ID。
num_nodes 是全局图中节点的数量。
num_edges 是全局图中的边数。
part-* 存储一个分区的数据。

当节点/边的ID被重新洗牌时，node_map 和 edge_map 包含了全局节点/边ID与分区本地节点/边ID之间的映射信息。对于异构图，node_map 和 edge_map 中的信息也可以用于计算节点类型和边类型。node_map 和 edge_map 中的数据格式如下：

{
    "node_type": [ [ part1_start, part1_end ],
                   [ part2_start, part2_end ],
                   ... ],
    ...
},

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

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

分区图包含额外的节点数据（“inner_node”）和边数据（“inner_edge”）：

“inner_node” 表示一个节点是否属于一个分区。
“inner_edge” 表示一条边是否属于一个分区。

节点和边的特征被分割并与每个图分区一起存储。分区中的所有节点/边特征都以DGL格式存储在一个文件中。节点/边特征存储在字典中，其中键是节点/边数据名称，值是一个张量。我们不存储HALO节点和边的特征。

在执行Metis分区时，我们可以在分区上施加一些约束。目前，它支持两种约束来平衡分区。默认情况下，Metis 总是尝试平衡每个分区中的节点数量。

balance_ntypes 平衡每个分区中不同类型节点的数量。
balance_edges 平衡每个分区中的边数。

为了平衡节点类型，用户需要传递一个包含N个元素的向量来指示每个节点的类型。N是输入图中节点的数量。

Parameters:

g (DGLGraph) – 要分区的输入图
graph_name (str) – 图的名称。该名称将用于构建 DistGraph()。
num_parts (int) – 分区数量
out_path (str) – 存储所有分区数据文件的路径。
num_hops (int, optional) – 我们在分区图结构上构建的HALO节点的跳数。默认值为1。
part_method (str, optional) – 分区方法。支持“random”和“metis”。默认值为“metis”。
balance_ntypes (tensor, optional) – 每个节点的节点类型。这是一个整数的一维数组。其值表示每个节点的节点类型。此参数由Metis分区使用。当指定此参数时，Metis算法将尝试将输入图划分为每个分区中每种节点类型的节点数量大致相同的分区。默认值为None，这意味着Metis将图分区以仅平衡节点数量。
balance_edges (bool) – 指示是否在每个分区中平衡边。此参数由Metis算法使用。
return_mapping (bool) – 指示是否返回打乱后的节点/边ID与原始节点/边ID之间的映射。
num_trainers_per_machine (int, 可选) – 每台机器上的训练器数量。如果不是1，整个图将首先被分区到每个训练器，即num_parts*num_trainers_per_machine个分区。每个节点的训练器ID将存储在节点特征‘trainer_id’中。然后，同一台机器上的训练器分区将被合并为一个更大的分区。最终的分区数量是num_part。
objtype (str, "cut" 或 "vol") – 将目标设置为边割最小化或通信量最小化。此参数由Metis算法使用。
graph_formats (str 或 list[str]) – 以指定格式保存分区。可以是 coo、csc 和 csr 的任意组合。如果未指定，则根据可用格式仅保存一种格式。如果有多种格式可用，选择优先级从高到低为 coo、csc、csr。
use_graphbolt (bool, optional) – 是否以GraphBolt格式保存分区。默认值：False。
kwargs (dict) – 用于将DGL分区转换为GraphBolt的其他关键字参数。

Returns:

Tensor 或 tensor 的字典，可选 – 如果 return_mapping=True，返回一个 1D tensor，表示同质图中打乱后的节点 ID 与原始节点 ID 之间的映射；返回一个 1D tensor 的字典，其键为节点类型，值为每个节点类型的打乱后的节点 ID 与原始节点 ID 之间的 1D tensor 映射。
Tensor 或 tensor 的字典，可选 – 如果 return_mapping=True，返回一个 1D tensor，表示同质图中打乱后的边 ID 与原始边 ID 之间的映射；返回一个 1D tensor 的字典，其键为边类型，值为每个边类型的打乱后的边 ID 与原始边 ID 之间的 1D tensor 映射。

示例

>>> dgl.distributed.partition_graph(g, 'test', 4, num_hops=1, part_method='metis',
...                                 out_path='output/',
...                                 balance_ntypes=g.ndata['train_mask'],
...                                 balance_edges=True)
>>> (
...     g, node_feats, edge_feats, gpb, graph_name, ntypes_list, etypes_list,
... ) = dgl.distributed.load_partition('output/test.json', 0)