| | 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 |
| |
|
