dgl.distributed.edge_split

dgl.distributed.edge_split(edges, partition_book=None, etype='_E', rank=None, force_even=True, edge_trainer_ids=None)[source]

划分边并返回本地 rank 的子集。

此函数根据分区簿划分输入边，并返回本地 rank 的一个边子集。此方法用于为分布式训练划分工作负载。

输入边可以存储为掩码向量。向量的长度与图中的边数相同；1 表示相应位置的边存在。

有两种划分边的策略。默认情况下，它以最大化数据局部性的方式划分边。也就是说，返回属于一个进程的所有边。如果 force_even 设置为 true，边将被平均划分，以便每个进程获得几乎相同数量的边。

当 force_even 为 True 时，如果图使用 Metis 划分并且节点/边 ID 已打乱，则数据局部性仍能得以保留。在这种情况下，为一个进程返回的大多数节点都是属于该进程的节点。如果节点/边 ID 未打乱，则无法保证数据局部性。

参数:

  • edges (1D tensor 或 DistTensor) – 表示输入边的布尔掩码向量。

  • partition_book (GraphPartitionBook, 可选) – 图分区簿

  • etype (str 或 (str, str, str), 可选) – 输入边的边类型。

  • rank (int, 可选) – 进程的 rank。如果未给出，则使用当前进程的 rank。

  • force_even (bool, 可选) – 强制平均划分边。

  • edge_trainer_ids (1D tensor 或 DistTensor, 可选) – 如果不为 None，则根据分配给每条边的 trainer ID 将边划分给同一机器上的 trainer。否则，随机划分。

返回值:

属于该 rank 的边 ID 向量。

返回类型:

1D 张量