File size: 2,147 Bytes
1239b39 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
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
|