dgl.segmented_knn_graph

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

根据k-近邻（KNN）从多组点构建多个图并返回。

与dgl.knn_graph()相比，这允许多个具有不同容量的点集。来自不同集合的点在x张量中连续存储。segs指定每个点集中的点数。该函数为每个点集构建一个KNN图，其中每个点的前驱是其通过欧几里得距离测量的k个最近邻居。DGL然后将所有KNN图组合成一个具有多个（\(len(segs)\)）连通组件的批量图。

Parameters:

  • x (Tensor) – 点的坐标/特征。必须是2D的。它可以在CPU或GPU上。

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

  • segs (list[int]) – 每个点集中的点数。segs中的数字必须加起来等于x中的行数。

  • algorithm (str, optional) –

    用于计算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, optional) – 用于计算点之间距离的距离度量。它可以是以下度量： * ‘euclidean’: 使用欧几里得距离（L2范数）\(\sqrt{\sum_{i} (x_{i} - y_{i})^{2}}\)。 * ‘cosine’: 使用余弦距离。 (默认: ‘euclidean’)

  • exclude_self (bool, optional) – If True, the output graph will not contain self loop edges, and each node will not be counted as one of its own k neighbors. If False, the output graph will contain self loop edges, and a node will be counted as one of its own k neighbors.

Returns:

批处理图。节点ID的顺序与x相同。

Return type:

DGLGraph

示例

以下示例使用 PyTorch 后端。

>>> import dgl
>>> import torch

在下面的示例中，第一个点集有三个点，第二个点集有四个点。

>>> # Features/coordinates of the first point set
>>> x1 = torch.tensor([[0.0, 0.5, 0.2],
...                    [0.1, 0.3, 0.2],
...                    [0.4, 0.2, 0.2]])
>>> # Features/coordinates of the second point set
>>> x2 = torch.tensor([[0.3, 0.2, 0.1],
...                    [0.5, 0.2, 0.3],
...                    [0.1, 0.1, 0.2],
...                    [0.6, 0.3, 0.3]])
>>> x = torch.cat([x1, x2], dim=0)
>>> segs = [x1.shape[0], x2.shape[0]]
>>> knn_g = dgl.segmented_knn_graph(x, 2, segs)
>>> knn_g.edges()
(tensor([0, 0, 1, 1, 1, 2, 3, 3, 4, 4, 5, 5, 6, 6]),
 tensor([0, 1, 0, 1, 2, 2, 3, 5, 4, 6, 3, 5, 4, 6]))