chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,384 @@
|
||||
"""Modules that transforms between graphs and between graph and tensors."""
|
||||
import torch.nn as nn
|
||||
|
||||
from ...transforms import knn_graph, radius_graph, segmented_knn_graph
|
||||
|
||||
|
||||
def pairwise_squared_distance(x):
|
||||
"""
|
||||
x : (n_samples, n_points, dims)
|
||||
return : (n_samples, n_points, n_points)
|
||||
"""
|
||||
x2s = (x * x).sum(-1, keepdim=True)
|
||||
return x2s + x2s.transpose(-1, -2) - 2 * x @ x.transpose(-1, -2)
|
||||
|
||||
|
||||
class KNNGraph(nn.Module):
|
||||
r"""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:
|
||||
|
||||
1. Compute an NxN matrix of pairwise distance for all points.
|
||||
2. Pick the k points with the smallest distance for each point as their k-nearest neighbors.
|
||||
3. Construct a graph with edges to each point as a node from its k-nearest neighbors.
|
||||
|
||||
The overall computational complexity is :math:`O(N^2(logN + D)`.
|
||||
|
||||
If a batch of point sets is provided, the point :math:`j` in point
|
||||
set :math:`i` is mapped to graph node ID: :math:`i \times M + j`, where
|
||||
:math:`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]))
|
||||
"""
|
||||
|
||||
def __init__(self, k):
|
||||
super(KNNGraph, self).__init__()
|
||||
self.k = k
|
||||
|
||||
# pylint: disable=invalid-name
|
||||
def forward(
|
||||
self,
|
||||
x,
|
||||
algorithm="bruteforce-blas",
|
||||
dist="euclidean",
|
||||
exclude_self=False,
|
||||
):
|
||||
r"""
|
||||
|
||||
Forward computation.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : Tensor
|
||||
:math:`(M, D)` or :math:`(N, M, D)` where :math:`N` means the
|
||||
number of point sets, :math:`M` means the number of points in
|
||||
each point set, and :math:`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 :math:`O(N^2)` memory complexity where
|
||||
:math:`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., :math:`O(Nk)` where :math:`N`
|
||||
is the number of points, :math:`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 <https://www.cs.princeton.edu/cass/papers/www11.pdf>`_. 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)
|
||||
:math:`\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
|
||||
-------
|
||||
DGLGraph
|
||||
A DGLGraph without features.
|
||||
"""
|
||||
return knn_graph(
|
||||
x, self.k, algorithm=algorithm, dist=dist, exclude_self=exclude_self
|
||||
)
|
||||
|
||||
|
||||
class SegmentedKNNGraph(nn.Module):
|
||||
r"""Layer that transforms one point set into a graph, or a batch of
|
||||
point sets with different number of points into a batched union of those graphs.
|
||||
|
||||
If a batch of point sets is provided, then the point :math:`j` in the point
|
||||
set :math:`i` is mapped to graph node ID:
|
||||
:math:`\sum_{p<i} |V_p| + j`, where :math:`|V_p|` means the number of
|
||||
points in the point set :math:`p`.
|
||||
|
||||
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 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]))
|
||||
>>>
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, k):
|
||||
super(SegmentedKNNGraph, self).__init__()
|
||||
self.k = k
|
||||
|
||||
# pylint: disable=invalid-name
|
||||
def forward(
|
||||
self,
|
||||
x,
|
||||
segs,
|
||||
algorithm="bruteforce-blas",
|
||||
dist="euclidean",
|
||||
exclude_self=False,
|
||||
):
|
||||
r"""Forward computation.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : Tensor
|
||||
:math:`(M, D)` where :math:`M` means the total number of points
|
||||
in all point sets, and :math:`D` means the size of features.
|
||||
segs : iterable of int
|
||||
:math:`(N)` integers where :math:`N` means the number of point
|
||||
sets. The number of elements must sum up to :math:`M`. And any
|
||||
:math:`N` should :math:`\ge k`
|
||||
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 :math:`O(N^2)` memory complexity where
|
||||
:math:`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., :math:`O(Nk)` where :math:`N`
|
||||
is the number of points, :math:`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 <https://www.cs.princeton.edu/cass/papers/www11.pdf>`_. 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)
|
||||
:math:`\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
|
||||
-------
|
||||
DGLGraph
|
||||
A batched DGLGraph without features.
|
||||
"""
|
||||
|
||||
return segmented_knn_graph(
|
||||
x,
|
||||
self.k,
|
||||
segs,
|
||||
algorithm=algorithm,
|
||||
dist=dist,
|
||||
exclude_self=exclude_self,
|
||||
)
|
||||
|
||||
|
||||
class RadiusGraph(nn.Module):
|
||||
r"""Layer that transforms one point set into a bidirected graph with
|
||||
neighbors within given distance.
|
||||
|
||||
The RadiusGraph is implemented in the following steps:
|
||||
|
||||
1. Compute an NxN matrix of pairwise distance for all points.
|
||||
2. Pick the points within distance to each point as their neighbors.
|
||||
3. Construct a graph with edges to each point as a node from its neighbors.
|
||||
|
||||
The nodes of the returned graph correspond to the points, where the neighbors
|
||||
of each point are within given distance.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
r : float
|
||||
Radius of the neighbors.
|
||||
p : float, optional
|
||||
Power parameter for the Minkowski metric. When :attr:`p = 1` it is the
|
||||
equivalent of Manhattan distance (L1 norm) and Euclidean distance
|
||||
(L2 norm) for :attr:`p = 2`.
|
||||
|
||||
(default: 2)
|
||||
self_loop : bool, optional
|
||||
Whether the radius graph will contain self-loops.
|
||||
|
||||
(default: False)
|
||||
compute_mode : str, optional
|
||||
``use_mm_for_euclid_dist_if_necessary`` - will use matrix multiplication
|
||||
approach to calculate euclidean distance (p = 2) if P > 25 or R > 25
|
||||
``use_mm_for_euclid_dist`` - will always use matrix multiplication
|
||||
approach to calculate euclidean distance (p = 2)
|
||||
``donot_use_mm_for_euclid_dist`` - will never use matrix multiplication
|
||||
approach to calculate euclidean distance (p = 2).
|
||||
|
||||
(default: donot_use_mm_for_euclid_dist)
|
||||
|
||||
Examples
|
||||
--------
|
||||
The following examples uses PyTorch backend.
|
||||
|
||||
>>> import dgl
|
||||
>>> from dgl.nn.pytorch.factory import RadiusGraph
|
||||
|
||||
>>> 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]])
|
||||
>>> rg = RadiusGraph(0.75)
|
||||
>>> g = rg(x) # Each node has neighbors within 0.75 distance
|
||||
>>> g.edges()
|
||||
(tensor([0, 1, 2, 2, 3, 3]), tensor([3, 2, 1, 3, 0, 2]))
|
||||
|
||||
When :attr:`get_distances` is True, forward pass returns the radius graph and
|
||||
distances for the corresponding edges.
|
||||
|
||||
>>> 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]])
|
||||
>>> rg = RadiusGraph(0.75)
|
||||
>>> g, dist = rg(x, get_distances=True)
|
||||
>>> g.edges()
|
||||
(tensor([0, 1, 2, 2, 3, 3]), tensor([3, 2, 1, 3, 0, 2]))
|
||||
>>> dist
|
||||
tensor([[0.7000],
|
||||
[0.6557],
|
||||
[0.6557],
|
||||
[0.2828],
|
||||
[0.7000],
|
||||
[0.2828]])
|
||||
"""
|
||||
|
||||
# pylint: disable=invalid-name
|
||||
def __init__(
|
||||
self,
|
||||
r,
|
||||
p=2,
|
||||
self_loop=False,
|
||||
compute_mode="donot_use_mm_for_euclid_dist",
|
||||
):
|
||||
super(RadiusGraph, self).__init__()
|
||||
self.r = r
|
||||
self.p = p
|
||||
self.self_loop = self_loop
|
||||
self.compute_mode = compute_mode
|
||||
|
||||
# pylint: disable=invalid-name
|
||||
def forward(self, x, get_distances=False):
|
||||
r"""
|
||||
Forward computation.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : Tensor
|
||||
The point coordinates. :math:`(N, D)` where :math:`N` means the
|
||||
number of points in the point set, and :math:`D` means the size of
|
||||
the features. It can be either on CPU or GPU. Device of the point
|
||||
coordinates specifies device of the radius graph.
|
||||
get_distances : bool, optional
|
||||
Whether to return the distances for the corresponding edges in the
|
||||
radius graph.
|
||||
|
||||
(default: False)
|
||||
|
||||
Returns
|
||||
-------
|
||||
DGLGraph
|
||||
The constructed graph. The node IDs are in the same order as :attr:`x`.
|
||||
torch.Tensor, optional
|
||||
The distances for the edges in the constructed graph. The distances
|
||||
are in the same order as edge IDs.
|
||||
"""
|
||||
return radius_graph(
|
||||
x, self.r, self.p, self.self_loop, self.compute_mode, get_distances
|
||||
)
|
||||
Reference in New Issue
Block a user