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)\)) 连通分量的批处理图。- 参数:
x (Tensor) – 点的坐标/特征。必须是二维的。可以在 CPU 或 GPU 上。
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’ 是一种近似方法，来自论文 用于通用相似性度量的高效 k 近邻图构建。此方法将在“邻居的邻居”中搜索最近邻候选节点。
(默认: ‘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相同。- 返回类型:
示例
以下示例使用 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]))