dgl.metis_partition

dgl.metis_partition(g, k, extra_cached_hops=0, reshuffle=False, balance_ntypes=None, balance_edges=False, mode='k-way')[source]

此 API 用于使用 Metis 分区算法对图进行分区。

Metis 将顶点分配给不同的分区。此 API 根据分配给分区的顶点及其入边构造子图。子图可能包含 HALO 节点，这些节点不属于子图所在的分区，但在固定跳数内与该分区中的节点相连。

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

  • balance_ntypes 用于平衡每个分区中不同类型的节点数量。

  • balance_edges 用于平衡每个分区中的边数量。

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

如果启用 reshuffle，函数会在分区前对输入图的节点 ID 和边 ID 进行重排。重排后，同一分区中的所有节点和边在输入图中会落在一个连续的 ID 范围内。分区后的子图包含节点数据 ‘orig_id’，用于存储原始输入图中的节点 ID。

分区后的子图存储在 DGLGraph 中。DGLGraph 包含 part_id 节点数据，指示节点所属的分区。子图不包含输入图中的节点/边数据。

参数:

  • g (DGLGraph) – 要分区的图

  • k (int) – 分区的数量。

  • extra_cached_hops (int) – HALO 节点可访问的跳数。

  • reshuffle (bool) – 是否重排节点，使同一分区的节点位于同一 ID 范围。

  • balance_ntypes (tensor) – 每个节点的节点类型

  • balance_edges (bool) – 指示是否平衡边数量。

  • mode (str, "k-way" or "recursive") – 使用多级递归二分法还是多级 k 向分区法。

返回值:

键是分区 ID，值是该分区的 DGLGraph。

返回类型:

包含 DGLGraphs 的字典