gather_points.py 1.55 KB
Newer Older
1
# Copyright (c) OpenMMLab. All rights reserved.
wuyuefeng's avatar
wuyuefeng committed
2
3
4
5
6
7
8
import torch
from torch.autograd import Function

from . import gather_points_ext


class GatherPoints(Function):
zhangwenwei's avatar
zhangwenwei committed
9
    """Gather Points.
wuyuefeng's avatar
wuyuefeng committed
10
11
12
13
14
15

    Gather points with given index.
    """

    @staticmethod
    def forward(ctx, features: torch.Tensor,
16
                indices: torch.Tensor) -> torch.Tensor:
wuyuefeng's avatar
wuyuefeng committed
17
18
19
20
        """forward.

        Args:
            features (Tensor): (B, C, N) features to gather.
21
            indices (Tensor): (B, M) where M is the number of points.
wuyuefeng's avatar
wuyuefeng committed
22
23
24
25
26

        Returns:
            Tensor: (B, C, M) where M is the number of points.
        """
        assert features.is_contiguous()
27
        assert indices.is_contiguous()
wuyuefeng's avatar
wuyuefeng committed
28

29
        B, npoint = indices.size()
wuyuefeng's avatar
wuyuefeng committed
30
31
32
33
        _, C, N = features.size()
        output = torch.cuda.FloatTensor(B, C, npoint)

        gather_points_ext.gather_points_wrapper(B, C, N, npoint, features,
34
                                                indices, output)
wuyuefeng's avatar
wuyuefeng committed
35

36
37
        ctx.for_backwards = (indices, C, N)
        ctx.mark_non_differentiable(indices)
wuyuefeng's avatar
wuyuefeng committed
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
        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()
        gather_points_ext.gather_points_grad_wrapper(B, C, N, npoint,
                                                     grad_out_data, idx,
                                                     grad_features.data)
        return grad_features, None


gather_points = GatherPoints.apply