【教程】DGL中的子图分区函数partition_graph讲解

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

实际上官方的函数解释中就已经非常详细了。

函数形式

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。

        分区数据存储在多个文件中，组织结构如下：

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": {
       "_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"中的数据格式如下：

{
    "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 张量映射。

使用示例

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