chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,12 @@
|
||||
"""The ``dgl.geometry`` package contains geometry operations:
|
||||
|
||||
* Farthest point sampling for point cloud sampling
|
||||
|
||||
* Neighbor matching module for graclus pooling
|
||||
|
||||
.. note::
|
||||
This package is experimental and the interfaces may be subject
|
||||
to changes in future releases.
|
||||
"""
|
||||
from .edge_coarsening import *
|
||||
from .fps import *
|
||||
@@ -0,0 +1,113 @@
|
||||
"""Python interfaces to DGL farthest point sampler."""
|
||||
import numpy as np
|
||||
|
||||
from .. import backend as F, ndarray as nd
|
||||
from .._ffi.base import DGLError
|
||||
from .._ffi.function import _init_api
|
||||
|
||||
|
||||
def _farthest_point_sampler(
|
||||
data, batch_size, sample_points, dist, start_idx, result
|
||||
):
|
||||
r"""Farthest Point Sampler
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tensor
|
||||
A tensor of shape (N, d) where N is the number of points and d is the dimension.
|
||||
batch_size : int
|
||||
The number of batches in the ``data``. N should be divisible by batch_size.
|
||||
sample_points : int
|
||||
The number of points to sample in each batch.
|
||||
dist : tensor
|
||||
Pre-allocated tensor of shape (N, ) for to-sample distance.
|
||||
start_idx : tensor of int
|
||||
Pre-allocated tensor of shape (batch_size, ) for the starting sample in each batch.
|
||||
result : tensor of int
|
||||
Pre-allocated tensor of shape (sample_points * batch_size, ) for the sampled index.
|
||||
|
||||
Returns
|
||||
-------
|
||||
No return value. The input variable ``result`` will be overwriten with sampled indices.
|
||||
|
||||
"""
|
||||
assert F.shape(data)[0] >= sample_points * batch_size
|
||||
assert F.shape(data)[0] % batch_size == 0
|
||||
|
||||
_CAPI_FarthestPointSampler(
|
||||
F.zerocopy_to_dgl_ndarray(data),
|
||||
batch_size,
|
||||
sample_points,
|
||||
F.zerocopy_to_dgl_ndarray(dist),
|
||||
F.zerocopy_to_dgl_ndarray(start_idx),
|
||||
F.zerocopy_to_dgl_ndarray(result),
|
||||
)
|
||||
|
||||
|
||||
def _neighbor_matching(
|
||||
graph_idx, num_nodes, edge_weights=None, relabel_idx=True
|
||||
):
|
||||
"""
|
||||
Description
|
||||
-----------
|
||||
The neighbor matching procedure of edge coarsening used in
|
||||
`Metis <http://cacs.usc.edu/education/cs653/Karypis-METIS-SIAMJSC98.pdf>`__
|
||||
and
|
||||
`Graclus <https://www.cs.utexas.edu/users/inderjit/public_papers/multilevel_pami.pdf>`__
|
||||
for homogeneous graph coarsening. This procedure keeps picking an unmarked
|
||||
vertex and matching it with one its unmarked neighbors (that maximizes its
|
||||
edge weight) until no match can be done.
|
||||
|
||||
If no edge weight is given, this procedure will randomly pick neighbor for each
|
||||
vertex.
|
||||
|
||||
The GPU implementation is based on `A GPU Algorithm for Greedy Graph Matching
|
||||
<http://www.staff.science.uu.nl/~bisse101/Articles/match12.pdf>`__
|
||||
|
||||
NOTE: The input graph must be bi-directed (undirected) graph. Call :obj:`dgl.to_bidirected`
|
||||
if you are not sure your graph is bi-directed.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
graph : HeteroGraphIndex
|
||||
The input homogeneous graph.
|
||||
num_nodes : int
|
||||
The number of nodes in this homogeneous graph.
|
||||
edge_weight : tensor, optional
|
||||
The edge weight tensor holding non-negative scalar weight for each edge.
|
||||
default: :obj:`None`
|
||||
relabel_idx : bool, optional
|
||||
If true, relabel resulting node labels to have consecutive node ids.
|
||||
default: :obj:`True`
|
||||
|
||||
Returns
|
||||
-------
|
||||
a 1-D tensor
|
||||
A vector with each element that indicates the cluster ID of a vertex.
|
||||
"""
|
||||
edge_weight_capi = nd.NULL["int64"]
|
||||
if edge_weights is not None:
|
||||
edge_weight_capi = F.zerocopy_to_dgl_ndarray(edge_weights)
|
||||
node_label = F.full_1d(
|
||||
num_nodes,
|
||||
-1,
|
||||
getattr(F, graph_idx.dtype),
|
||||
F.to_backend_ctx(graph_idx.ctx),
|
||||
)
|
||||
node_label_capi = F.zerocopy_to_dgl_ndarray_for_write(node_label)
|
||||
_CAPI_NeighborMatching(graph_idx, edge_weight_capi, node_label_capi)
|
||||
if F.reduce_sum(node_label < 0).item() != 0:
|
||||
raise DGLError("Find unmatched node")
|
||||
|
||||
# reorder node id
|
||||
# TODO: actually we can add `return_inverse` option for `unique`
|
||||
# function in backend for efficiency.
|
||||
if relabel_idx:
|
||||
node_label_np = F.zerocopy_to_numpy(node_label)
|
||||
_, node_label_np = np.unique(node_label_np, return_inverse=True)
|
||||
return F.tensor(node_label_np)
|
||||
else:
|
||||
return node_label
|
||||
|
||||
|
||||
_init_api("dgl.geometry", __name__)
|
||||
@@ -0,0 +1,64 @@
|
||||
"""Edge coarsening procedure used in Metis and Graclus, for pytorch"""
|
||||
# pylint: disable=no-member, invalid-name, W0613
|
||||
from .. import remove_self_loop
|
||||
from .capi import _neighbor_matching
|
||||
|
||||
__all__ = ["neighbor_matching"]
|
||||
|
||||
|
||||
def neighbor_matching(graph, e_weights=None, relabel_idx=True):
|
||||
r"""
|
||||
Description
|
||||
-----------
|
||||
The neighbor matching procedure of edge coarsening in
|
||||
`Metis <http://cacs.usc.edu/education/cs653/Karypis-METIS-SIAMJSC98.pdf>`__
|
||||
and
|
||||
`Graclus <https://www.cs.utexas.edu/users/inderjit/public_papers/multilevel_pami.pdf>`__
|
||||
for homogeneous graph coarsening. This procedure keeps picking an unmarked
|
||||
vertex and matching it with one its unmarked neighbors (that maximizes its
|
||||
edge weight) until no match can be done.
|
||||
|
||||
If no edge weight is given, this procedure will randomly pick neighbor for each
|
||||
vertex.
|
||||
|
||||
The GPU implementation is based on `A GPU Algorithm for Greedy Graph Matching
|
||||
<http://www.staff.science.uu.nl/~bisse101/Articles/match12.pdf>`__
|
||||
|
||||
NOTE: The input graph must be bi-directed (undirected) graph. Call :obj:`dgl.to_bidirected`
|
||||
if you are not sure your graph is bi-directed.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
graph : DGLGraph
|
||||
The input homogeneous graph.
|
||||
edge_weight : torch.Tensor, optional
|
||||
The edge weight tensor holding non-negative scalar weight for each edge.
|
||||
default: :obj:`None`
|
||||
relabel_idx : bool, optional
|
||||
If true, relabel resulting node labels to have consecutive node ids.
|
||||
default: :obj:`True`
|
||||
|
||||
Examples
|
||||
--------
|
||||
The following example uses PyTorch backend.
|
||||
|
||||
>>> import torch, dgl
|
||||
>>> from dgl.geometry import neighbor_matching
|
||||
>>>
|
||||
>>> g = dgl.graph(([0, 1, 1, 2], [1, 0, 2, 1]))
|
||||
>>> res = neighbor_matching(g)
|
||||
tensor([0, 1, 1])
|
||||
"""
|
||||
assert (
|
||||
graph.is_homogeneous
|
||||
), "The graph used in graph node matching must be homogeneous"
|
||||
if e_weights is not None:
|
||||
graph.edata["e_weights"] = e_weights
|
||||
graph = remove_self_loop(graph)
|
||||
e_weights = graph.edata["e_weights"]
|
||||
graph.edata.pop("e_weights")
|
||||
else:
|
||||
graph = remove_self_loop(graph)
|
||||
return _neighbor_matching(
|
||||
graph._graph, graph.num_nodes(), e_weights, relabel_idx
|
||||
)
|
||||
@@ -0,0 +1,65 @@
|
||||
"""Farthest Point Sampler for pytorch Geometry package"""
|
||||
# pylint: disable=no-member, invalid-name
|
||||
|
||||
from .. import backend as F
|
||||
from ..base import DGLError
|
||||
from .capi import _farthest_point_sampler
|
||||
|
||||
__all__ = ["farthest_point_sampler"]
|
||||
|
||||
|
||||
def farthest_point_sampler(pos, npoints, start_idx=None):
|
||||
"""Farthest Point Sampler without the need to compute all pairs of distance.
|
||||
|
||||
In each batch, the algorithm starts with the sample index specified by ``start_idx``.
|
||||
Then for each point, we maintain the minimum to-sample distance.
|
||||
Finally, we pick the point with the maximum such distance.
|
||||
This process will be repeated for ``sample_points`` - 1 times.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
pos : tensor
|
||||
The positional tensor of shape (B, N, C)
|
||||
npoints : int
|
||||
The number of points to sample in each batch.
|
||||
start_idx : int, optional
|
||||
If given, appoint the index of the starting point,
|
||||
otherwise randomly select a point as the start point.
|
||||
(default: None)
|
||||
|
||||
Returns
|
||||
-------
|
||||
tensor of shape (B, npoints)
|
||||
The sampled indices in each batch.
|
||||
|
||||
Examples
|
||||
--------
|
||||
The following exmaple uses PyTorch backend.
|
||||
|
||||
>>> import torch
|
||||
>>> from dgl.geometry import farthest_point_sampler
|
||||
>>> x = torch.rand((2, 10, 3))
|
||||
>>> point_idx = farthest_point_sampler(x, 2)
|
||||
>>> print(point_idx)
|
||||
tensor([[5, 6],
|
||||
[7, 8]])
|
||||
"""
|
||||
ctx = F.context(pos)
|
||||
B, N, C = pos.shape
|
||||
pos = pos.reshape(-1, C)
|
||||
dist = F.zeros((B * N), dtype=pos.dtype, ctx=ctx)
|
||||
if start_idx is None:
|
||||
start_idx = F.randint(
|
||||
shape=(B,), dtype=F.int64, ctx=ctx, low=0, high=N - 1
|
||||
)
|
||||
else:
|
||||
if start_idx >= N or start_idx < 0:
|
||||
raise DGLError(
|
||||
"Invalid start_idx, expected 0 <= start_idx < {}, got {}".format(
|
||||
N, start_idx
|
||||
)
|
||||
)
|
||||
start_idx = F.full_1d(B, start_idx, dtype=F.int64, ctx=ctx)
|
||||
result = F.zeros((npoints * B), dtype=F.int64, ctx=ctx)
|
||||
_farthest_point_sampler(pos, B, npoints, dist, start_idx, result)
|
||||
return result.reshape(B, npoints)
|
||||
Reference in New Issue
Block a user