dgl.distributed.node_split

dgl.distributed.node_split(nodes, partition_book=None, ntype='_N', rank=None, force_even=True, node_trainer_ids=None)[source]

分割节点并返回本地等级的子集。

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

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

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

当 force_even 为 True 时，如果使用 Metis 对图进行分区并且节点/边 ID 被打乱，数据局部性仍然得以保留。 在这种情况下，返回给进程的大多数节点是属于该进程的节点。如果节点/边 ID 没有被打乱，则不能保证数据局部性。

Parameters:

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

  • partition_book (GraphPartitionBook, optional) – 图分区书

  • ntype (str, optional) – 输入节点的节点类型。

  • rank (int, 可选) – 进程的排名。如果未提供，则使用当前进程的排名。

  • force_even (bool, optional) – 强制节点均匀分割。

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

Returns:

属于该等级的节点ID向量。

Return type:

一维张量