|
from typing import Tuple
|
|
|
|
import torch
|
|
from torch.autograd import Function
|
|
|
|
from ..utils import ext_loader
|
|
|
|
ext_module = ext_loader.load_ext(
|
|
'_ext', ['three_interpolate_forward', 'three_interpolate_backward'])
|
|
|
|
|
|
class ThreeInterpolate(Function):
|
|
"""Performs weighted linear interpolation on 3 features.
|
|
|
|
Please refer to `Paper of PointNet++ <https://arxiv.org/abs/1706.02413>`_
|
|
for more details.
|
|
"""
|
|
|
|
@staticmethod
|
|
def forward(ctx, features: torch.Tensor, indices: torch.Tensor,
|
|
weight: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Args:
|
|
features (Tensor): (B, C, M) Features descriptors to be
|
|
interpolated
|
|
indices (Tensor): (B, n, 3) index three nearest neighbors
|
|
of the target features in features
|
|
weight (Tensor): (B, n, 3) weights of interpolation
|
|
|
|
Returns:
|
|
Tensor: (B, C, N) tensor of the interpolated features
|
|
"""
|
|
assert features.is_contiguous()
|
|
assert indices.is_contiguous()
|
|
assert weight.is_contiguous()
|
|
|
|
B, c, m = features.size()
|
|
n = indices.size(1)
|
|
ctx.three_interpolate_for_backward = (indices, weight, m)
|
|
output = torch.cuda.FloatTensor(B, c, n)
|
|
|
|
ext_module.three_interpolate_forward(
|
|
features, indices, weight, output, b=B, c=c, m=m, n=n)
|
|
return output
|
|
|
|
@staticmethod
|
|
def backward(
|
|
ctx, grad_out: torch.Tensor
|
|
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
"""
|
|
Args:
|
|
grad_out (Tensor): (B, C, N) tensor with gradients of outputs
|
|
|
|
Returns:
|
|
Tensor: (B, C, M) tensor with gradients of features
|
|
"""
|
|
idx, weight, m = ctx.three_interpolate_for_backward
|
|
B, c, n = grad_out.size()
|
|
|
|
grad_features = torch.cuda.FloatTensor(B, c, m).zero_()
|
|
grad_out_data = grad_out.data.contiguous()
|
|
|
|
ext_module.three_interpolate_backward(
|
|
grad_out_data, idx, weight, grad_features.data, b=B, c=c, n=n, m=m)
|
|
return grad_features, None, None
|
|
|
|
|
|
three_interpolate = ThreeInterpolate.apply
|
|
|
