KNNGraphļ
- class dgl.nn.pytorch.factory.KNNGraph(k)[source]ļ
Bases:
Module
Layer that transforms one point set into a graph, or a batch of point sets with the same number of points into a batched union of those graphs.
The KNNGraph is implemented in the following steps:
Compute an NxN matrix of pairwise distance for all points.
Pick the k points with the smallest distance for each point as their k-nearest neighbors.
Construct a graph with edges to each point as a node from its k-nearest neighbors.
The overall computational complexity is \(O(N^2(logN + D)\).
If a batch of point sets is provided, the point \(j\) in point set \(i\) is mapped to graph node ID: \(i \times M + j\), where \(M\) is the number of nodes in each point set.
The predecessors of each node are the k-nearest neighbors of the corresponding point.
- Parameters:
k (int) ā The number of neighbors.
Notes
The nearest neighbors found for a node include the node itself.
Examples
The following example uses PyTorch backend.
>>> import torch >>> from dgl.nn.pytorch.factory import KNNGraph >>> >>> kg = KNNGraph(2) >>> x = torch.tensor([[0,1], [1,2], [1,3], [100, 101], [101, 102], [50, 50]]) >>> g = kg(x) >>> print(g.edges()) (tensor([0, 1, 1, 1, 2, 2, 2, 3, 3, 4, 4, 5]), tensor([0, 0, 1, 2, 1, 2, 5, 3, 4, 3, 4, 5]))
- forward(x, algorithm='bruteforce-blas', dist='euclidean', exclude_self=False)[source]ļ
Forward computation.
- Parameters:
x (Tensor) ā \((M, D)\) or \((N, M, D)\) where \(N\) means the number of point sets, \(M\) means the number of points in each point set, and \(D\) means the size of features.
algorithm (str, optional) ā
Algorithm used to compute the k-nearest neighbors.
ābruteforce-blasā will first compute the distance matrix using BLAS matrix multiplication operation provided by backend frameworks. Then use topk algorithm to get k-nearest neighbors. This method is fast when the point set is small but has \(O(N^2)\) memory complexity where \(N\) is the number of points.
ābruteforceā will compute distances pair by pair and directly select the k-nearest neighbors during distance computation. This method is slower than ābruteforce-blasā but has less memory overhead (i.e., \(O(Nk)\) where \(N\) is the number of points, \(k\) is the number of nearest neighbors per node) since we do not need to store all distances.
ābruteforce-sharememā (CUDA only) is similar to ābruteforceā but use shared memory in CUDA devices for buffer. This method is faster than ābruteforceā when the dimension of input points is not large. This method is only available on CUDA device.
ākd-treeā will use the kd-tree algorithm (CPU only). This method is suitable for low-dimensional data (e.g. 3D point clouds)
ānn-descentā is a approximate approach from paper Efficient k-nearest neighbor graph construction for generic similarity measures. This method will search for nearest neighbor candidates in āneighborsā neighborsā.
(default: ābruteforce-blasā)
dist (str, optional) ā
The distance metric used to compute distance between points. It can be the following metrics: * āeuclideanā: Use Euclidean distance (L2 norm)
\(\sqrt{\sum_{i} (x_{i} - y_{i})^{2}}\).
ācosineā: Use cosine distance.
(default: ā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:
A DGLGraph without features.
- Return type: