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