add pointops lib for eval; add example slurm train script

Author: HaoyiZhu

libs/pointops/__init__.py

@@ -0,0 +1 @@
+from .functions import *

libs/pointops/functions/__init__.py

@@ -0,0 +1,14 @@
+from .aggregation import aggregation
+from .attention import attention_fusion_step, attention_relation_step
+from .grouping import grouping, grouping2
+from .interpolation import interpolation, interpolation2
+from .query import ball_query, knn_query, random_ball_query
+from .sampling import farthest_point_sampling
+from .subtraction import subtraction
+from .utils import (
+    ball_query_and_group,
+    batch2offset,
+    knn_query_and_group,
+    offset2batch,
+    query_and_group,
+)

libs/pointops/functions/aggregation.py

@@ -0,0 +1,56 @@
+import torch
+from pointops._C import aggregation_backward_cuda, aggregation_forward_cuda
+from torch.autograd import Function
+
+
+class Aggregation(Function):
+    @staticmethod
+    def forward(ctx, input, position, weight, idx):
+        """
+        input: input: (n, c), position: (n, nsample, c), weight : (n, nsample, c'), idx: (n, nsample)
+        output: (n, c)
+        """
+        assert (
+            input.is_contiguous()
+            and position.is_contiguous()
+            and weight.is_contiguous()
+        )
+        n, nsample, c = position.shape
+        w_c = weight.shape[-1]
+        output = torch.cuda.FloatTensor(n, c).zero_()
+        aggregation_forward_cuda(
+            n, nsample, c, w_c, input, position, weight, idx, output
+        )
+        ctx.save_for_backward(input, position, weight, idx)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        """
+        input: grad_out: (n, c)
+        output: grad_input: (n, c), grad_position: (n, nsample, c), grad_weight : (n, nsample, c')
+        """
+        input, position, weight, idx = ctx.saved_tensors
+        n, nsample, c = position.shape
+        w_c = weight.shape[-1]
+        grad_input = torch.cuda.FloatTensor(n, c).zero_()
+        grad_position = torch.cuda.FloatTensor(n, nsample, c).zero_()
+        grad_weight = torch.cuda.FloatTensor(n, nsample, w_c).zero_()
+        aggregation_backward_cuda(
+            n,
+            nsample,
+            c,
+            w_c,
+            input,
+            position,
+            weight,
+            idx,
+            grad_output,
+            grad_input,
+            grad_position,
+            grad_weight,
+        )
+        return grad_input, grad_position, grad_weight, None
+
+
+aggregation = Aggregation.apply

libs/pointops/functions/attention.py

@@ -0,0 +1,119 @@
+import torch
+from pointops._C import (
+    attention_fusion_step_backward_cuda,
+    attention_fusion_step_forward_cuda,
+    attention_relation_step_backward_cuda,
+    attention_relation_step_forward_cuda,
+)
+from torch.autograd import Function
+
+
+class AttentionRelationStep(Function):
+    @staticmethod
+    def forward(ctx, query, key, weight, index_target, index_refer):
+        """
+        input - query: (n, g, c), key: (n, g, c), weight: (c)  1_c for scatter attention,
+                index_target: (m), index_refer: (m)
+        output - relation: (M, g)
+        """
+
+        assert (
+            query.is_contiguous()
+            and key.is_contiguous()
+            and index_target.is_contiguous()
+            and index_refer.is_contiguous()
+            and weight.is_contiguous()
+        )
+
+        assert index_target.shape[0] == index_refer.shape[0]
+
+        _, g, c = query.shape
+        m = index_target.shape[0]
+        output = torch.cuda.FloatTensor(m, g).zero_()
+        attention_relation_step_forward_cuda(
+            m, g, c, query, key, weight, index_target.int(), index_refer.int(), output
+        )
+        ctx.save_for_backward(query, key, weight, index_target, index_refer)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        query, key, weight, index_target, index_refer = ctx.saved_tensors
+        n, g, c = query.shape
+        m = index_target.shape[0]
+        grad_query = torch.cuda.FloatTensor(n, g, c).zero_()
+        grad_key = torch.cuda.FloatTensor(n, g, c).zero_()
+        grad_weight = torch.cuda.FloatTensor(c).zero_()
+        attention_relation_step_backward_cuda(
+            m,
+            g,
+            c,
+            query,
+            grad_query,
+            key,
+            grad_key,
+            weight,
+            grad_weight,
+            index_target.int(),
+            index_refer.int(),
+            grad_output,
+        )
+        return grad_query, grad_key, None, None, None
+
+
+class AttentionFusionStep(Function):
+    @staticmethod
+    def forward(ctx, weight, value, index_target, index_refer):
+        """
+        input - weight: (m, g), value: (n, g, c)
+                index_target: (m), index_value: (m)
+        output - output: (n, g, c)
+        """
+
+        assert (
+            weight.is_contiguous()
+            and value.is_contiguous()
+            and index_target.is_contiguous()
+            and index_refer.is_contiguous()
+            and weight.is_contiguous()
+        )
+
+        assert index_target.shape[0] == index_refer.shape[0]
+
+        n, g, c = value.shape
+        m = index_refer.shape[0]
+        output = torch.cuda.FloatTensor(n, g, c).zero_()
+        attention_fusion_step_forward_cuda(
+            m, g, c, weight, value, index_target.int(), index_refer.int(), output
+        )
+        ctx.save_for_backward(weight, value, index_target, index_refer)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        """
+        input: grad_output: (n, g, c)
+        output: grad_weight: (m, g), grad_value: (n, g, c), none, none
+        """
+        weight, value, index_target, index_refer = ctx.saved_tensors
+        n, g, c = value.shape
+        m = index_target.shape[0]
+        grad_weight = torch.cuda.FloatTensor(m, g).zero_()
+        grad_value = torch.cuda.FloatTensor(n, g, c).zero_()
+        attention_fusion_step_backward_cuda(
+            m,
+            g,
+            c,
+            weight,
+            grad_weight,
+            value,
+            grad_value,
+            index_target.int(),
+            index_refer.int(),
+            grad_output,
+        )
+        return grad_weight, grad_value, None, None
+
+
+attention_relation_step = AttentionRelationStep.apply
+attention_fusion_step = AttentionFusionStep.apply

libs/pointops/functions/grouping.py

@@ -0,0 +1,62 @@
+import torch
+from pointops._C import grouping_backward_cuda, grouping_forward_cuda
+from torch.autograd import Function
+
+
+class Grouping(Function):
+    @staticmethod
+    def forward(ctx, input, idx):
+        """
+        input: input: (n, c), idx : (m, nsample)
+        output: (m, nsample, c)
+        """
+        assert input.is_contiguous() and idx.is_contiguous()
+        m, nsample, n, c = idx.shape[0], idx.shape[1], input.shape[0], input.shape[1]
+        output = torch.cuda.FloatTensor(m, nsample, c)
+        grouping_forward_cuda(m, nsample, c, input, idx, output)
+        ctx.n = n
+        ctx.save_for_backward(idx)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        """
+        input: grad_out: (m, c, nsample)
+        output: (n, c), None
+        """
+        n = ctx.n
+        (idx,) = ctx.saved_tensors
+        m, nsample, c = grad_output.shape
+        grad_input = torch.cuda.FloatTensor(n, c).zero_()
+        grouping_backward_cuda(m, nsample, c, grad_output, idx, grad_input)
+        return grad_input, None
+
+
+def grouping(idx, feat, xyz, new_xyz=None, with_xyz=False):
+    if new_xyz is None:
+        new_xyz = xyz
+    assert xyz.is_contiguous() and feat.is_contiguous()
+    m, nsample, c = idx.shape[0], idx.shape[1], feat.shape[1]
+    xyz = torch.cat([xyz, torch.zeros([1, 3]).to(xyz.device)], dim=0)
+    feat = torch.cat([feat, torch.zeros([1, c]).to(feat.device)], dim=0)
+    grouped_feat = feat[idx.view(-1).long(), :].view(
+        m, nsample, c
+    )  # (m, num_sample, c)
+
+    if with_xyz:
+        assert new_xyz.is_contiguous()
+        mask = torch.sign(idx + 1)
+        grouped_xyz = xyz[idx.view(-1).long(), :].view(
+            m, nsample, 3
+        ) - new_xyz.unsqueeze(
+            1
+        )  # (m, num_sample, 3)
+        grouped_xyz = torch.einsum(
+            "n s c, n s -> n s c", grouped_xyz, mask
+        )  # (m, num_sample, 3)
+        return torch.cat((grouped_xyz, grouped_feat), -1)
+    else:
+        return grouped_feat
+
+
+grouping2 = Grouping.apply

libs/pointops/functions/interpolation.py

@@ -0,0 +1,59 @@
+import torch
+from pointops._C import interpolation_backward_cuda, interpolation_forward_cuda
+from torch.autograd import Function
+
+from .query import knn_query
+
+
+def interpolation(xyz, new_xyz, feat, offset, new_offset, k=3):
+    """
+    input: coords: (m, 3), new_xyz: (n, 3), color: (m, c), offset: (b), new_offset: (b)
+    output: (n, c)
+    """
+    assert xyz.is_contiguous() and new_xyz.is_contiguous() and feat.is_contiguous()
+    idx, dist = knn_query(k, xyz, offset, new_xyz, new_offset)  # (n, 3), (n, 3)
+    dist_recip = 1.0 / (dist + 1e-8)  # (n, 3)
+    norm = torch.sum(dist_recip, dim=1, keepdim=True)
+    weight = dist_recip / norm  # (n, 3)
+
+    new_feat = torch.cuda.FloatTensor(new_xyz.shape[0], feat.shape[1]).zero_()
+    for i in range(k):
+        new_feat += feat[idx[:, i].long(), :] * weight[:, i].unsqueeze(-1)
+    return new_feat
+
+
+class Interpolation(Function):
+    @staticmethod
+    def forward(ctx, xyz, new_xyz, input, offset, new_offset, k=3):
+        """
+        input: coords: (m, 3), new_xyz: (n, 3), input: (m, c), offset: (b), new_offset: (b)
+        output: (n, c)
+        """
+        assert xyz.is_contiguous() and new_xyz.is_contiguous() and input.is_contiguous()
+        idx, dist = knn_query(k, xyz, offset, new_xyz, new_offset)  # (n, k), (n, k)
+        dist_recip = 1.0 / (dist + 1e-8)  # (n, k)
+        norm = torch.sum(dist_recip, dim=1, keepdim=True)
+        weight = dist_recip / norm  # (n, k)
+
+        n, c, m = new_xyz.shape[0], input.shape[1], input.shape[0]
+        output = torch.cuda.FloatTensor(n, c).zero_()
+        interpolation_forward_cuda(n, c, k, input, idx, weight, output)
+        ctx.m, ctx.k = m, k
+        ctx.save_for_backward(idx, weight)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        """
+        input: coords: (m, 3), new_xyz: (n, 3), input: (m, c), offset: (b), new_offset: (b)
+        output: (n, c)
+        """
+        m, k = ctx.m, ctx.k
+        idx, weight = ctx.saved_tensors
+        n, c = grad_output.shape
+        grad_input = torch.cuda.FloatTensor(m, c).zero_()
+        interpolation_backward_cuda(n, c, k, grad_output, idx, weight, grad_input)
+        return None, None, grad_input, None, None, None
+
+
+interpolation2 = Interpolation.apply