dgl.distributed.edge_split

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

分割边并返回本地排名的子集。

此函数根据分区书拆分输入边，并返回本地等级的边子集。此方法用于分配分布式训练的工作负载。

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

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

When force_even is True, the data locality is still preserved if a graph is partitioned with Metis and the node/edge IDs are shuffled. In this case, majority of the nodes returned for a process are the ones that belong to the process. If node/edge IDs are not shuffled, data locality is not guaranteed.

Parameters:

  • edges (1D tensor 或 DistTensor) – 一个布尔掩码向量，用于指示输入边。

  • partition_book (GraphPartitionBook, optional) – The graph partition book

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

  • rank (int, optional) – The rank of a process. If not given, the rank of the current process is used.

  • force_even (bool, optional) – 强制边缘均匀分割。

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

Returns:

属于该等级的边的ID向量。

Return type:

一维张量