dmlc--dgl
Tong He
3d47693b1f
* working framework without actual algorithm logic * rename * fix * fps passes compilation * correct algorithm * add cuda implementation * update random start * before refactor * pass compilation but cuda not working * working * code working, will add docstring * add mxnet support * update docstring * update doc and test * cpplint * cpcplint * pylint * temporary fix * fix for win64 * fix unitetest * fix * fix * remove comment * move to geometry package * remove redundant include * add docstrings and comments * add proof * add validity check
48 行
1.5 KiB
Python
48 行
1.5 KiB
Python
"""Farthest Point Sampler for mxnet Geometry package"""
|
|
#pylint: disable=no-member, invalid-name
|
|
|
|
from mxnet import nd
|
|
from mxnet.gluon import nn
|
|
import numpy as np
|
|
|
|
from ..capi import farthest_point_sampler
|
|
|
|
class FarthestPointSampler(nn.Block):
|
|
"""Farthest Point Sampler
|
|
|
|
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
|
|
----------
|
|
npoints : int
|
|
The number of points to sample in each batch.
|
|
"""
|
|
def __init__(self, npoints):
|
|
super(FarthestPointSampler, self).__init__()
|
|
self.npoints = npoints
|
|
|
|
def forward(self, pos):
|
|
r"""Memory allocation and sampling
|
|
|
|
Parameters
|
|
----------
|
|
pos : tensor
|
|
The positional tensor of shape (B, N, C)
|
|
|
|
Returns
|
|
-------
|
|
tensor of shape (B, self.npoints)
|
|
The sampled indices in each batch.
|
|
"""
|
|
ctx = pos.context
|
|
B, N, C = pos.shape
|
|
pos = pos.reshape(-1, C)
|
|
dist = nd.zeros((B * N), dtype=pos.dtype, ctx=ctx)
|
|
start_idx = nd.random.randint(0, N - 1, (B, ), dtype=np.int, ctx=ctx)
|
|
result = nd.zeros((self.npoints * B), dtype=np.int, ctx=ctx)
|
|
farthest_point_sampler(pos, B, self.npoints, dist, start_idx, result)
|
|
return result.reshape(B, self.npoints)
|