dgl.knn_graph

dgl.knn_graph(x, k, algorithm='bruteforce-blas', dist='euclidean', exclude_self=False)[source]

根据 k-最近邻 (KNN) 构建图并返回。

该函数将点集的坐标/特征转换为有向齐次图。点集的坐标被指定为一个矩阵，其中行对应于点，列对应于坐标/特征维度。

返回图的节点对应于点，其中每个点的前驱节点是根据所选距离测量的 k 个最近邻。

如果 x 是一个 3D tensor，则每个子矩阵将被转换为一个独立的图。然后 DGL 将这些图组合成一个由多个 (\(shape(x)[0]\)) 连通分量组成的大型批处理图。

有关完整的基准测试结果，请参阅基准测试。

参数:

  • x (Tensor) –

    点坐标。它可以在 CPU 或 GPU 上。

    • 如果是 2D，x[i] 对应于 KNN 图中的第 i 个节点。

    • 如果是 3D，x[i] 对应于第 i 个 KNN 图，而 x[i][j] 对应于第 i 个 KNN 图中的第 j 个节点。

  • k (int) – 每个节点的最近邻数量。

  • algorithm (str, 可选) –

    用于计算 k-最近邻的算法。

    • ‘bruteforce-blas’ 将首先使用后端框架提供的 BLAS 矩阵乘法操作计算距离矩阵。然后使用 topk 算法获取 k-最近邻。当点集较小时，此方法速度很快，但内存复杂度为 \(O(N^2)\)，其中 \(N\) 是点数。

    • ‘bruteforce’ 将逐对计算距离，并在距离计算期间直接选择 k-最近邻。此方法比 ‘bruteforce-blas’ 慢，但内存开销较小（即 \(O(Nk)\)，其中 \(N\) 是点数，\(k\) 是每个节点的最近邻数量），因为我们无需存储所有距离。

    • ‘bruteforce-sharemem’（仅限 CUDA）类似于 ‘bruteforce’，但在 CUDA 设备中使用共享内存作为缓冲区。当输入点维度不大时，此方法比 ‘bruteforce’ 快。此方法仅在 CUDA 设备上可用。

    • ‘kd-tree’ 将使用 kd-tree 算法（仅限 CPU）。此方法适用于低维数据（例如 3D 点云）

    • ‘nn-descent’ 是论文 Efficient k-nearest neighbor graph construction for generic similarity measures 中的一种近似方法。此方法将在“邻居的邻居”中搜索最近邻候选者。

    (默认值: ‘bruteforce-blas’)

  • dist (str, 可选) – 用于计算点之间距离的距离度量。可以是以下度量之一：* ‘euclidean’: 使用欧几里得距离 (L2 范数) \(\sqrt{\sum_{i} (x_{i} - y_{i})^{2}}\)。* ‘cosine’: 使用余弦距离。(默认值: ‘euclidean’)

  • exclude_self (bool, 可选) – 如果为 True，则输出图将不包含自环边，并且每个节点不会被视为其自身的 k 个邻居之一。如果为 False，则输出图将包含自环边，并且一个节点将被视为其自身的 k 个邻居之一。

返回值:

构建的图。节点 ID 的顺序与 x 相同。

返回类型:

DGLGraph

示例

以下示例使用 PyTorch 后端。

>>> import dgl
>>> import torch

当 x 是一个 2D tensor 时，构建一个单一的 KNN 图。

>>> x = torch.tensor([[0.0, 0.0, 1.0],
...                   [1.0, 0.5, 0.5],
...                   [0.5, 0.2, 0.2],
...                   [0.3, 0.2, 0.4]])
>>> knn_g = dgl.knn_graph(x, 2)  # Each node has two predecessors
>>> knn_g.edges()
(tensor([0, 1, 2, 2, 2, 3, 3, 3]), tensor([0, 1, 1, 2, 3, 0, 2, 3]))

当 x 是一个 3D tensor 时，DGL 构建多个 KNN 图，然后将它们组合成一个由多个连通分量组成的图。

>>> x1 = torch.tensor([[0.0, 0.0, 1.0],
...                    [1.0, 0.5, 0.5],
...                    [0.5, 0.2, 0.2],
...                    [0.3, 0.2, 0.4]])
>>> x2 = torch.tensor([[0.0, 1.0, 1.0],
...                    [0.3, 0.3, 0.3],
...                    [0.4, 0.4, 1.0],
...                    [0.3, 0.8, 0.2]])
>>> x = torch.stack([x1, x2], dim=0)
>>> knn_g = dgl.knn_graph(x, 2)  # Each node has two predecessors
>>> knn_g.edges()
(tensor([0, 1, 2, 2, 2, 3, 3, 3, 4, 5, 5, 5, 6, 6, 7, 7]),
 tensor([0, 1, 1, 2, 3, 0, 2, 3, 4, 5, 6, 7, 4, 6, 5, 7]))