dgl.graphbolt.unique_and_compact_csc_formats

dgl.graphbolt.unique_and_compact_csc_formats(csc_formats: Tuple[Tensor, Tensor] | Dict[str, Tuple[Tensor, Tensor]], unique_dst_nodes: Tensor | Dict[str, Tensor], rank: int = 0, world_size: int = 1, async_op: bool = False)[源码]

压缩 CSC 格式并返回（按类型分类的）唯一节点。使用协作小批量（Cooperative Minibatching）时，rank 和 world_size 参数相关，协作小批量最初在 `Deep Graph Library PR#4337`__ 中提出，后来在图神经网络中的协作小批量中首次完整描述。GPU 之间的协作消除了由于在执行 GNN 小批量处理时，种子节点的 k 跳采样邻居存在重叠而在 GPU 之间重复执行的工作。

当 world_size 大于 1 时，给定的 ID 会在可用的 ranks 之间进行划分。保证与给定 rank 对应的 ID 会出现在其他 ranks 的 ID 之前。为此，划分后的 ID 会按给定 rank 向后旋转，以便 ID 的顺序为：[rank, rank + 1, world_size, 0, …, rank - 1]。这仅支持 Volta 及后续代次的 NVIDIA GPU。

参数：

csc_formats (Union[CSCFormatBase, Dict(str, CSCFormatBase)]) – 表示源-目的边的 CSC 格式。- 如果 csc_formats 是 CSCFormatBase：表示图是同构的。此外，其中的 indptr 和 indice 应该是 torch.Tensor，表示 CSC 格式的源-目的对。内部的 ID 是同构 ID。- 如果 csc_formats 是 Dict[str, CSCFormatBase]：键应该是边类型，值应该是 CSC 格式的节点对。内部的 ID 是异构 ID。
unique_dst_nodes (torch.Tensor or Dict[str, torch.Tensor]) – 节点对中所有目的节点的唯一节点。- 如果 unique_dst_nodes 是一个 tensor：表示图是同构的。- 如果 unique_dst_nodes 是一个字典：键是节点类型，值是对应的节点。内部的 ID 是异构 ID。
rank (int) – 当前进程的 rank。
world_size (int) – 进程数量。
async_op (bool) – 布尔值，指示调用是否为异步。如果是，可以通过调用返回的 future 的 wait 方法来获取结果。

返回：

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

返回类型：

Tuple[unique_nodes, csc_formats, unique_nodes_offsets]

示例

>>> import dgl.graphbolt as gb
>>> N1 = torch.LongTensor([1, 2, 2])
>>> N2 = torch.LongTensor([5, 5, 6])
>>> unique_dst = {
...     "n1": torch.LongTensor([1, 2]),
...     "n2": torch.LongTensor([5, 6])}
>>> csc_formats = {
...     "n1:e1:n2": gb.CSCFormatBase(indptr=torch.tensor([0, 2, 3]),indices=N1),
...     "n2:e2:n1": gb.CSCFormatBase(indptr=torch.tensor([0, 1, 3]),indices=N2)}
>>> unique_nodes, compacted_csc_formats, _ = gb.unique_and_compact_csc_formats(
...     csc_formats, unique_dst
... )
>>> print(unique_nodes)
{'n1': tensor([1, 2]), 'n2': tensor([5, 6])}
>>> print(compacted_csc_formats)
{"n1:e1:n2": CSCFormatBase(indptr=torch.tensor([0, 2, 3]),
                           indices=torch.tensor([0, 1, 1])),
 "n2:e2:n1": CSCFormatBase(indptr=torch.tensor([0, 1, 3]),
                           indices=torch.Longtensor([0, 0, 1]))}