SegmentedKNNGraph

class dgl.nn.pytorch.factory.SegmentedKNNGraph(k)[source]

Bases: Module

一个将点集转换为图，或将一批具有不同点数的点集转换为这些图的批处理联合的层。

如果提供了一批点集，则点集 $i$ 中的点 $j$ 被映射到图节点 ID： $\sum_{p < i} | V_{p} | + j$ ，其中 $| V_{p} |$ 表示点集 $p$ 中的点数。

每个节点的前驱是对应点的 k 个最近邻。

参数：: k (int) – 邻居的数量。

注意事项

为节点找到的最近邻包括节点本身。

示例

以下示例使用 PyTorch 后端。

>>> import torch
>>> from dgl.nn.pytorch.factory import SegmentedKNNGraph
>>>
>>> kg = SegmentedKNNGraph(2)
>>> x = torch.tensor([[0,1],
...                   [1,2],
...                   [1,3],
...                   [100, 101],
...                   [101, 102],
...                   [50, 50],
...                   [24,25],
...                   [25,24]])
>>> g = kg(x, [3,3,2])
>>> print(g.edges())
(tensor([0, 1, 1, 1, 2, 2, 3, 3, 3, 4, 4, 5, 6, 6, 7, 7]),
 tensor([0, 0, 1, 2, 1, 2, 3, 4, 5, 3, 4, 5, 6, 7, 6, 7]))
>>>

forward(x, segs, algorithm='bruteforce-blas', dist='euclidean', exclude_self=False)[source]

前向计算。

参数：

x (Tensor) – $(M, D)$ ，其中 $M$ 表示所有点集中的点总数， $D$ 表示特征的大小。
segs (iterable of int) – $(N)$ 个整数，其中 $N$ 表示点集的数量。元素数量的总和必须等于 $M$ 。并且任何一个 $N$ 都应 $\geq k$
algorithm (str, optional) –
用于计算 k 个最近邻的算法。
- ‘bruteforce-blas’ 将首先使用后端框架提供的 BLAS 矩阵乘法操作计算距离矩阵。然后使用 topk 算法获取 k 个最近邻。当点集较小时，此方法速度很快，但内存复杂度为 $O (N^{2})$ ，其中 $N$ 是点数。
- ‘bruteforce’ 将成对计算距离，并在距离计算过程中直接选择 k 个最近邻。此方法比 ‘bruteforce-blas’ 慢，但内存开销更小（即 $O (N k)$ ，其中 $N$ 是点数， $k$ 是每个节点的最近邻数量），因为我们不需要存储所有距离。
- ‘bruteforce-sharemem’（仅限 CUDA）类似于 ‘bruteforce’，但在 CUDA 设备中使用共享内存作为缓冲区。当输入点维度不大时，此方法比 ‘bruteforce’ 快。此方法仅适用于 CUDA 设备。
- ‘kd-tree’ 将使用 kd-tree 算法（仅限 CPU）。此方法适用于低维数据（例如 3D 点云）
- ‘nn-descent’ 是论文用于通用相似度度量的有效 k 最近邻图构建中的近似方法。此方法将在“邻居的邻居”中搜索最近邻候选者。
（默认值：‘bruteforce-blas’）
dist (str, optional) –
用于计算点之间距离的度量。可以是以下度量之一：

* ‘euclidean’：使用欧氏距离 (L2 范数).
- ’cosine’：使用余弦距离。
（默认值：‘euclidean’）
exclude_self (bool, optional) – 如果为 True，则输出图将不包含自循环边，并且每个节点不会被视为其自身的 k 个邻居之一。如果为 False，则输出图将包含自循环边，并且节点将被视为其自身的 k 个邻居之一。

返回：

一个不带特征的批处理 DGLGraph。

返回类型：

DGLGraph