dgl.graphbolt.unique_and_compact

dgl.graphbolt.unique_and_compact(nodes: List[Tensor] | Dict[str, List[Tensor]], rank: int = 0, world_size: int = 1, async_op: bool = False)[source]

压缩节点张量列表。 rank 和 world_size 参数在使用协同 MiniBatching 时相关，协同 MiniBatching 最初在 `Deep Graph Library PR#4337<https://github.com/dmlc/dgl/pull/4337>`__ 中提出，后来在 图神经网络中的协同 MiniBatching 中首次完整描述。GPU 之间的协同消除了在执行 GNN MiniBatching 时，由于采样的种子节点 k 跳邻居重叠而在不同 GPU 上产生的重复工作。

当 world_size 大于 1 时，给定的 ID 会在可用的 rank 之间进行划分。给定 rank 对应的 ID 保证位于其他 rank 的 ID 之前。为此，通过给定 rank 将划分的 ID 向后旋转，使得 ID 的顺序为：[rank, rank + 1, world_size, 0, …, rank - 1]。此功能仅支持 Volta 及更新一代的 NVIDIA GPU。

参数：

  • nodes (List[torch.Tensor] or Dict[str, List[torch.Tensor]]) – 用于压缩的节点张量列表。 unique_and_compact 将按类型进行。 - 如果 nodes 是张量列表：所有张量将一起进行去重和压缩，通常用于同构图。 - 如果 nodes 是字典列表：键应为节点类型，值应为对应的节点，去重和压缩将按类型进行，通常用于异构图。

  • rank (int) – 当前进程的 rank。

  • world_size (int) – 进程数量。

  • async_op (bool) – 布尔值，指示调用是否异步。如果异步，可以通过对返回的 future 调用 wait 获取结果。

返回值：

输入中所有节点的唯一节点（按类型）。以及压缩后的节点列表，其中内部的 ID 被替换为压缩后的节点 ID。“压缩后的节点列表”表示输入节点列表中的节点 ID 已被映射的节点 ID 替换，其中每种类型的节点都映射到一个从 0 到 N 的连续 ID 空间。唯一节点偏移量张量划分了 unique_nodes 张量。其大小为 world_size + 1，且 unique_nodes[offsets[i]: offsets[i + 1]] 属于 rank (rank + i) % world_size。

返回值类型：

Tuple[unique_nodes, compacted_node_list, unique_nodes_offsets]