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 是输入图中的节点数量。

参数:

g (DGLGraph) – 要划分的输入图
graph_name (str) – 图的名称。该名称将用于构建 DistGraph()。
num_parts (int) – 分区数量
out_path (str) – 存储所有分区数据文件的路径。
num_hops (int, 可选) – 在分区图结构上构建的 HALO 节点跳数。默认值为 1。
part_method (str, 可选) – 分区方法。支持“random”和“metis”。默认值为“metis”。
balance_ntypes (张量, 可选) – 每个节点的节点类型。这是一个整数组成的一维数组。其值指示每个节点的节点类型。此参数由 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") – 设置目标为最小化边割 (edge-cut) 或最小化通信量 (communication volume)。此参数由 Metis 算法使用。
graph_formats (str 或 list[str]) – 以指定格式保存分区。可以是 coo、csc 和 csr 的任意组合。如果未指定，则仅根据可用格式保存一种格式。如果多种格式可用，优先级从高到低依次是 coo、csc、csr。
use_graphbolt (bool, 可选) – 是否以 GraphBolt 格式保存分区。默认值：False。
kwargs (dict) – 将 DGL 分区转换为 GraphBolt 的其他关键字参数。

返回值:

张量或张量字典, 可选 – 如果 return_mapping=True，则对于同构图，返回一个指示打乱后的节点 ID 与原始节点 ID 之间映射关系的一维张量；对于异构图，返回一个一维张量字典，其中键是节点类型，值是一个一维张量，表示每种节点类型的打乱后的节点 ID 与原始节点 ID 之间的映射。
张量或张量字典, 可选 – 如果 return_mapping=True，则对于同构图，返回一个指示打乱后的边 ID 与原始边 ID 之间映射关系的一维张量；对于异构图，返回一个一维张量字典，其中键是边类型，值是一个一维张量，表示每种边类型的打乱后的边 ID 与原始边 ID 之间的映射。

示例

>>> 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)