|
import torch
|
|
from torch.autograd import Function
|
|
|
|
from ..utils import ext_loader
|
|
|
|
ext_module = ext_loader.load_ext(
|
|
'_ext', ['gather_points_forward', 'gather_points_backward'])
|
|
|
|
|
|
class GatherPoints(Function):
|
|
"""Gather points with given index."""
|
|
|
|
@staticmethod
|
|
def forward(ctx, features: torch.Tensor,
|
|
indices: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Args:
|
|
features (Tensor): (B, C, N) features to gather.
|
|
indices (Tensor): (B, M) where M is the number of points.
|
|
|
|
Returns:
|
|
Tensor: (B, C, M) where M is the number of points.
|
|
"""
|
|
assert features.is_contiguous()
|
|
assert indices.is_contiguous()
|
|
|
|
B, npoint = indices.size()
|
|
_, C, N = features.size()
|
|
output = torch.cuda.FloatTensor(B, C, npoint)
|
|
|
|
ext_module.gather_points_forward(
|
|
features, indices, output, b=B, c=C, n=N, npoints=npoint)
|
|
|
|
ctx.for_backwards = (indices, C, N)
|
|
if torch.__version__ != 'parrots':
|
|
ctx.mark_non_differentiable(indices)
|
|
return output
|
|
|
|
@staticmethod
|
|
def backward(ctx, grad_out):
|
|
idx, C, N = ctx.for_backwards
|
|
B, npoint = idx.size()
|
|
|
|
grad_features = torch.cuda.FloatTensor(B, C, N).zero_()
|
|
grad_out_data = grad_out.data.contiguous()
|
|
ext_module.gather_points_backward(
|
|
grad_out_data,
|
|
idx,
|
|
grad_features.data,
|
|
b=B,
|
|
c=C,
|
|
n=N,
|
|
npoints=npoint)
|
|
return grad_features, None
|
|
|
|
|
|
gather_points = GatherPoints.apply
|
|
|
