7.4 高级图划分

本章涵盖图划分的一些高级主题。

METIS 划分算法

METIS 是一种最先进的图划分算法，可以生成跨分区的边数量最少的划分，这使得它适用于分布式消息传递，因为网络通信量与跨分区的边数量成比例。DGL 在其 dgl.distributed.partition_graph() API 中集成了 METIS 作为默认的划分算法。

输出格式

无论使用哪种划分算法，划分结果都存储在数据文件中，组织方式如下：

data_root_dir/
  |-- graph_name.json       # partition configuration file in JSON
  |-- 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
  |
  |-- ...                   # data for other partitions

当分发到集群时，元数据 JSON 应该复制到所有机器，而 partX 文件夹应该相应地分派。

DGL 提供 dgl.distributed.load_partition() 函数来加载单个分区进行检查。

>>> import dgl
>>> # load partition 0
>>> part_data = dgl.distributed.load_partition('data_root_dir/graph_name.json', 0)
>>> g, nfeat, efeat, partition_book, graph_name, ntypes, etypes = part_data  # unpack
>>> print(g)
Graph(num_nodes=966043, num_edges=34270118,
      ndata_schemes={'orig_id': Scheme(shape=(), dtype=torch.int64),
                     'part_id': Scheme(shape=(), dtype=torch.int64),
                     '_ID': Scheme(shape=(), dtype=torch.int64),
                     'inner_node': Scheme(shape=(), dtype=torch.int32)}
      edata_schemes={'_ID': Scheme(shape=(), dtype=torch.int64),
                     'inner_edge': Scheme(shape=(), dtype=torch.int8),
                     'orig_id': Scheme(shape=(), dtype=torch.int64)})

如`ID mapping`_部分所述，每个分区都包含作为 ndata 或 edata 保存的辅助信息，例如原始节点/边 ID、分区 ID 等。每个分区不仅保存其拥有的节点/边，还包括与分区相邻的节点/边（称为 HALO 节点/边）。inner_node 和 inner_edge 指示节点/边是否真正属于该分区（值为 True）还是 HALO 节点/边（值为 False）。

通过 load_partition() 函数一次性加载所有数据。用户可以使用 dgl.distributed.load_partition_feats() 和 dgl.distributed.load_partition_book() API 分别加载特征或分区簿。

并行 METIS 划分

对于需要并行预处理的大规模图，DGL 支持 ParMETIS 作为其中一种划分算法选择。

注意

因为 ParMETIS 不支持异构图，用户需要在运行 ParMETIS 前后进行 ID 转换。请查看第 7.5 章异构图深入解析获取解释。

注意

请确保 ParMETIS 的输入图不包含重复边（或平行边）和自环边。

ParMETIS 安装

ParMETIS 需要 METIS 和 GKLib。请按照此处的说明编译和安装 GKLib。对于 METIS 的编译和安装，请按照下面的说明使用 GIT 克隆 METIS 并编译使其支持 int64。

git clone https://github.com/KarypisLab/METIS.git
make config shared=1 cc=gcc prefix=~/local i64=1
make install

目前，我们需要手动编译和安装 ParMETIS。我们克隆 ParMETIS 的 DGL 分支如下：

git clone --branch dgl https://github.com/KarypisLab/ParMETIS.git

然后编译并安装 ParMETIS。

make config cc=mpicc prefix=~/local
make install

在运行 ParMETIS 之前，我们需要设置两个环境变量：PATH 和 LD_LIBRARY_PATH。

export PATH=$PATH:$HOME/local/bin
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$HOME/local/lib/

输入格式

注意

作为先决条件，请阅读 guide-distributed-hetero 章以了解 DGL 如何组织异构图进行分布式训练。

ParMETIS 的输入图存储在三个文件中，名称如下：xxx_nodes.txt、xxx_edges.txt 和 xxx_stats.txt，其中 xxx 是图的名称。

xxx_nodes.txt 中的每一行存储一个节点的信息。行 ID 也是节点的齐次 ID，例如，第 0 行对应节点 0；第 1 行对应节点 1 等。每行的格式如下：

<node_type_id> <node_weight_list> <type_wise_node_id>

所有字段都以空格分隔

<node_type_id> 是从 0 开始的整数。每个节点类型都被映射到一个整数。对于齐次图，其值始终为 0。
<node_weight_list> 是整数（以空格分隔），表示 ParMETIS 用于平衡图分区的节点权重。对于齐次图，此列表只有一个整数；而对于包含 \(T\) 个节点类型的异构图，此列表应包含 \(T\) 个整数。如果节点属于节点类型 \(t\)，则除第 \(t\) 个整数外，所有整数都为零；第 \(t\) 个整数是该节点的权重。ParMETIS 将尝试平衡每个分区的总节点权重。对于异构图，它将尝试将相同类型的节点分布到所有分区。推荐的节点权重是 1（用于平衡每个分区中的节点数量）或节点度数（用于平衡每个分区中的边数量）。
<type_wise_node_id> 是一个整数，表示节点在其自身类型中的 ID。

下面展示了一个包含两种节点类型的异构图的节点文件示例。节点类型 0 有三个节点；节点类型 1 有四个节点。它使用两个节点权重来确保 ParMETIS 生成的分区中类型 0 和类型 1 的节点数量大致相同。

类似地，xxx_edges.txt 中的每一行存储一条边信息。行 ID 也是边的齐次 ID，例如，第 0 行对应边 0；第 1 行对应边 1 等。每行的格式如下：

<src_node_id> <dst_node_id> <type_wise_edge_id> <edge_type_id>

所有字段都以空格分隔

<src_node_id> 是源节点的齐次 ID。
<dst_node_id> 是目标节点的齐次 ID。
<type_wise_edge_id> 是该边类型中的边 ID。
<edge_type_id> 是从 0 开始的整数。每个边类型都被映射到一个整数。对于齐次图，其值始终为 0。

xxx_stats.txt 存储图的一些基本统计信息。它只有一行，包含三个字段，以空格分隔：

<num_nodes> <num_edges> <total_node_weights>

num_nodes 存储节点总数，不区分节点类型。
num_edges 存储边总数，不区分边类型。
total_node_weights 存储节点文件中的节点权重数量。

运行 ParMETIS 并输出格式

ParMETIS 包含一个名为 pm_dglpart 的命令，它从调用 pm_dglpart 的机器加载存储在三个文件中的图，将数据分发到集群中的所有机器，然后调用 ParMETIS 对图进行划分。完成后，它会为每个分区生成三个文件：p<part_id>-xxx_nodes.txt、p<part_id>-xxx_edges.txt 和 p<part_id>-xxx_stats.txt。

注意

ParMETIS 在划分过程中会重新分配节点的 ID。ID 重新分配后，分区中的节点会被分配连续的 ID；此外，相同类型的节点也会被分配连续的 ID。

p<part_id>-xxx_nodes.txt 存储分区中的节点数据。每行代表一个节点，包含以下字段：

<node_id> <node_type_id> <node_weight_list> <type_wise_node_id>

<node_id> 是 ID 重新分配后的齐次节点 ID。
<node_type_id> 是节点类型 ID。
<node_weight_list> 是 ParMETIS 使用的节点权重（从输入文件复制）。
<type_wise_node_id> 是一个整数，表示节点在其自身类型中的 ID。

p<part_id>-xxx_edges.txt 存储分区中的边数据。每行代表一条边，包含以下字段：

<src_id> <dst_id> <orig_src_id> <orig_dst_id> <type_wise_edge_id> <edge_type_id>

<src_id> 是 ID 重新分配后源节点的齐次 ID。
<dst_id> 是 ID 重新分配后目标节点的齐次 ID。
<orig_src_id> 是输入图中的源节点的齐次 ID。
<orig_dst_id> 是输入图中的目标节点的齐次 ID。
<type_wise_edge_id> 是该边类型中的边 ID。
<edge_type_id> 是边类型 ID。

调用 pm_dglpart 时，三个输入文件：xxx_nodes.txt、xxx_edges.txt 和 xxx_stats.txt 应位于运行 pm_dglpart 的目录中。以下命令运行四个 ParMETIS 进程将名为 xxx 的图划分为八个分区（每个进程处理两个分区）。

mpirun -np 4 pm_dglpart xxx 2

ParMETIS 的输出文件需要转换为划分分配格式，以便运行后续的预处理步骤。