Create new file

Intra_MLP.py

@@ -0,0 +1,90 @@
+import torch
+import numpy
+#from transformer import Local_Attention,Transformer_1
+# codes of this function are borrowed from https://github.com/yanx27/Pointnet_Pointnet2_pytorch/blob/master/models/pointnet2_utils.py
+def index_points(device, points, idx):
+    """
+
+    Input:
+        points: input points data, [B, N, C]
+        idx: sample index data, [B, S]
+    Return:
+        new_points:, indexed points data, [B, S, C]
+    """
+    B = points.shape[0]
+    view_shape = list(idx.shape)
+    view_shape[1:] = [1] * (len(view_shape) - 1)
+    repeat_shape = list(idx.shape)
+    repeat_shape[0] = 1
+    # batch_indices = torch.arange(B, dtype=torch.long).to(device).view(view_shape).repeat(repeat_shape)
+    batch_indices = torch.arange(B, dtype=torch.long).to(device).view(view_shape).repeat(repeat_shape)
+    new_points = points[batch_indices, idx, :]
+    return new_points
+
+def knn_l2(device, net, k, u):
+    '''
+    Input:
+        k: int32, number of k in k-nn search
+        net: (batch_size, npoint, c) float32 array, points
+        u: int32, block size
+    Output:
+        idx: (batch_size, npoint, k) int32 array, indices to input points
+    '''
+    INF = 1e8
+    batch_size = net.size(0)
+    npoint = net.size(1)
+    n_channel = net.size(2)
+
+    square = torch.pow(torch.norm(net, dim=2,keepdim=True),2)
+
+    def u_block(batch_size, npoint, u):
+        block = numpy.zeros([batch_size, npoint, npoint])
+        n = npoint // u
+        for i in range(n):
+            block[:, (i*u):(i*u+u), (i*u):(i*u+u)] = numpy.ones([batch_size, u, u]) * (-INF)
+        return block
+
+    # minus_distance = 2 * torch.matmul(net, net.transpose(2,1)) - square - square.transpose(2,1) + torch.Tensor(u_block(batch_size, npoint, u)).to(device)
+    minus_distance = 2 * torch.matmul(net, net.transpose(2,1)) - square - square.transpose(2,1) + torch.Tensor(u_block(batch_size, npoint, u)).to(device)
+    _, indices = torch.topk(minus_distance, k, largest=True, sorted=False)
+    
+    return indices
+
+if __name__ == '__main__':
+
+        bs,gs,k=5,5,4
+        
+        A=torch.rand(bs*gs,512,14,14).cuda()
+        net=Transformer_1(512,4,4,782).cuda()
+        Y=net(A)
+        print(Y.shape)
+        exit(0)
+        feature_map_size=A.shape[-1]
+        point = A.permute(0,2,1,3,4).reshape(A.size(0), A.size(1)*A.shape[-1]*A.shape[-2], -1)
+        point = point.permute(0,2,1)
+        X=point
+        print(point.shape)
+        idx = knn_l2(0, point, 4, 1)
+        #print(idx)
+        
+        feat=idx
+        new_point = index_points(0, point,idx)
+
+        group_point = new_point.permute(0, 3, 2, 1)
+        print(group_point.shape)
+        _1,_2,_3,_4=group_point.shape
+        X=X.permute(0,2,1)
+        print(X.shape)
+        #torch.cat([group_point.reshape(_1*_2,k,_4),X.reshape(_1*_2,1,_4)],dim=1).permute(0,2,1)
+        attn_map=X.reshape(_1*_2,1,_4)@torch.cat([group_point.reshape(_1*_2,k,_4),X.reshape(_1*_2,1,_4)],dim=1).permute(0,2,1)
+        V=torch.cat([group_point.reshape(_1*_2,k,_4),X.reshape(_1*_2,1,_4)],dim=1)
+        print(attn_map.shape)
+        Y=attn_map@V
+        Y=Y.reshape(_1,_2,_4)
+        
+        #group_point = torch.max(group_point, 2)[0]  # [B, D', S]
+        group_point=Y
+        print(group_point.shape)
+        
+        intra_mask = group_point.view(bs,gs, group_point.size(2), feature_map_size, feature_map_size)
+        print(intra_mask.shape)