[Model] Add random_start option for Farthest Point Sampler (dmlc#2755)

* add random_start option for FPS * change doc * change from random_start to start_idx * change error condition * change error msg Co-authored-by: Minjie Wang <[email protected]>
feixian15 · Mar 22, 2021 · 6999f88 · 6999f88
1 parent bcffdb8
commit 6999f88
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 4 deletions.
diff --git a/python/dgl/geometry/mxnet/fps.py b/python/dgl/geometry/mxnet/fps.py
@@ -5,6 +5,7 @@
 from mxnet.gluon import nn
 import numpy as np
 
+from ...base import DGLError
 from ..capi import farthest_point_sampler
 
 class FarthestPointSampler(nn.Block):
@@ -24,13 +25,17 @@ def __init__(self, npoints):
         super(FarthestPointSampler, self).__init__()
         self.npoints = npoints
 
-    def forward(self, pos):
+    def forward(self, pos, start_idx=None):
         r"""Memory allocation and sampling
 
         Parameters
         ----------
         pos : tensor
             The positional tensor of shape (B, N, C)
+        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
         -------
@@ -41,7 +46,13 @@ def forward(self, pos):
         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)
+        if start_idx is None:
+            start_idx = nd.random.randint(0, N - 1, (B, ), dtype=np.int, ctx=ctx)
+        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 = nd.full((B, ), start_idx, 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)
diff --git a/python/dgl/geometry/pytorch/fps.py b/python/dgl/geometry/pytorch/fps.py
@@ -4,6 +4,7 @@
 import torch as th
 from torch import nn
 
+from ...base import DGLError
 from ..capi import farthest_point_sampler
 
 class FarthestPointSampler(nn.Module):
@@ -23,13 +24,17 @@ def __init__(self, npoints):
         super(FarthestPointSampler, self).__init__()
         self.npoints = npoints
 
-    def forward(self, pos):
+    def forward(self, pos, start_idx=None):
         r"""Memory allocation and sampling
 
         Parameters
         ----------
         pos : tensor
             The positional tensor of shape (B, N, C)
+        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
         -------
@@ -40,7 +45,13 @@ def forward(self, pos):
         B, N, C = pos.shape
         pos = pos.reshape(-1, C)
         dist = th.zeros((B * N), dtype=pos.dtype, device=device)
-        start_idx = th.randint(0, N - 1, (B, ), dtype=th.long, device=device)
+        if start_idx is None:
+            start_idx = th.randint(0, N - 1, (B, ), dtype=th.long, device=device)
+        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 = th.full((B, ), start_idx, dtype=th.long, device=device)
         result = th.zeros((self.npoints * B), dtype=th.long, device=device)
         farthest_point_sampler(pos, B, self.npoints, dist, start_idx, result)
         return result.reshape(B, self.npoints)