dgl.distributed.node_split

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

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

此函数根据分区簿（partition book）分割输入节点，并返回本地 rank 的节点子集。此方法用于划分分布式训练的工作负载。

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

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

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

参数：

  • nodes (1D tensor or DistTensor) – 指示输入节点的布尔掩码向量。

  • partition_book (GraphPartitionBook, optional) – 图分区簿。

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

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

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

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

返回：

属于该 rank 的节点 ID 向量。

返回类型：

1D-tensor