def _compute_weights(self, X):
#X = check_array(X)
#self.nbrs_ = NearestNeighbors(n_neighbors=self.n_neighbors,
# algorithm=self.neighbors_algorithm)
#self.nbrs_.fit(X)
points_number, *para = X.size()
knn = KNN(k=self.n_neighbors,
transpose_mode=True) #points,neighbours..
dist, col = knn(X.unsqueeze(0), X.unsqueeze(0))
#knn_pc = torch.cat([X[0,col[0,:,i],:].unsqueeze(-1) for i in range(self.n_neighbours)],dim=-1).permute(0,2,1)
W = torch.zeros([points_number, points_number],
dtype=X.dtype,
device=X.device)
for i in range(5):
W[:, col[0, :, i]] = 1
if self.weight == 'adjacency':
W = W #kneighbors_graph(self.nbrs_, self.n_neighbors,
# mode='connectivity', include_self=True)
elif self.weight == 'heat':
W = W #kneighbors_graph(self.nbrs_, self.n_neighbors,
# mode='distance', include_self=True)
W = torch.exp(-W**2 / self.weight_width**2)
else:
raise ValueError("Unrecognized Weight")
# symmetrize the matrix
# TODO: make this more efficient & keep sparse output
#W = W.toarray()
p = torch.cat([W.unsqueeze(-1), W.t().unsqueeze(-1)], dim=-1)
W = torch.max(p, dim=-1, keepdim=False)[0]
#print(W.size())
return W
def knn(x, k: int):
"""
inputs:
- x: b x npoints1 x num_dims (partical_cloud)
- k: int (the number of neighbor)
outputs:
- idx: int (neighbor_idx)
"""
# x : (batch_size, feature_dim, num_points)
# Retrieve nearest neighbor indices
if torch.cuda.is_available():
from knn_cuda import KNN
ref = x.transpose(2, 1).contiguous() # (batch_size, num_points, feature_dim)
query = ref
_, idx = KNN(k=k, transpose_mode=True)(ref, query)
else:
inner = -2 * torch.matmul(x.transpose(2, 1), x)
xx = torch.sum(x ** 2, dim=1, keepdim=True)
pairwise_distance = -xx - inner - xx.transpose(2, 1)
idx = pairwise_distance.topk(k=k, dim=-1)[1] # (batch_size, num_points, k)
return idx
def __init__(self,
in_channels,
coord_channels,
out_channels,
K,
mlp=[32, 64],
point_pooling=False,
attentive_aggregation=False):
super(ContinuousConvolution, self).__init__()
self.K = K
self.C_in = in_channels
self.C_coord = coord_channels
self.C_out = out_channels
self.C_hid = mlp
self.num_output_features = self.C_out
self.knn = KNN(k=self.K, transpose_mode=True)
self.kernel = MLP_conv1x1(channels=[self.C_in + self.C_coord] +
self.C_hid + [self.C_out])
self.point_pooling = point_pooling
if self.point_pooling:
self.num_output_features += self.C_in
self.attentive_aggregation = attentive_aggregation
if self.attentive_aggregation:
self.aggr_mlp = MLP_conv1x1(channels=[self.K, 1])
self.num_output_features += self.C_out
def forward(self, coords, features):
r"""
Forward pass
Parameters
----------
coords: torch.Tensor, shape (B, N, 3)
coordinates of the point cloud
features: torch.Tensor, shape (B, d_in, N, 1)
features of the point cloud
Returns
-------
torch.Tensor, shape (B, 2*d_out, N, 1)
"""
knn1 = KNN(k=1, transpose_mode=True)
knn_output = knn1(coords.cuda().contiguous(),
coords.cuda().contiguous())
x = self.mlp1(features)
x = self.lse1(coords, x, knn_output)
x = self.pool1(x)
x = self.lse2(coords, x, knn_output)
x = self.pool2(x)
return self.lrelu(self.mlp2(x) + self.shortcut(features))
def _project_and_match(self, x: torch.Tensor,
simp: torch.Tensor) -> (torch.Tensor, torch.Tensor):
match = None
proj = None
# Projected points
if self.training:
if not self.skip_projection:
proj = self.project(point_cloud=x, query_cloud=simp)
else:
proj = simp
# Matched points
else: # Inference
# Retrieve nearest neighbor indices
_, idx = KNN(1, transpose_mode=False)(x.contiguous(),
simp.contiguous())
"""Notice that we detach the tensors and do computations in numpy,
and then convert back to Tensors.
This should have no effect as the network is in eval() mode
and should require no gradients.
"""
# Convert to numpy arrays in B x N x 3 format. we assume 'bcn' format.
x = x.permute(0, 2, 1).cpu().detach().numpy()
idx = idx.cpu().detach().numpy()
idx = np.squeeze(idx, axis=1)
z = sputils.nn_matching(x,
idx,
self.num_out_points,
complete_fps=self.complete_fps)
# Matched points are in B x N x 3 format.
match = torch.tensor(z, dtype=torch.float32).cuda()
# Change to output shapes
if self.output_shape == "bnc":
simp = simp.permute(0, 2, 1)
if proj is not None:
proj = proj.permute(0, 2, 1)
elif self.output_shape == "bcn" and match is not None:
match = match.permute(0, 2, 1)
match = match.contiguous()
# Assert contiguous tensors
simp = simp.contiguous()
if proj is not None:
proj = proj.contiguous()
out = proj
else:
match = match.contiguous()
out = match
if self.debug:
return simp, proj, match
else:
return simp, out
def sample_and_group_cuda(npoint,
k,
xyz,
points,
cat_xyz_feature=True,
fps_only=False):
"""
Input:
npoint:
k:
xyz: input points position data, [B, N, 3]
points: input points data, [B, C, N]
Return:
new_xyz: sampled points position data, [B, 3, npoint]
new_points: sampled points data, [B, C+C_xyz, npoint, k]
grouped_xyz_norm: sampled relative points position data, [B, 3, npoint, k]
"""
k = min(npoint, k)
knn = KNN(k=k, transpose_mode=True)
B, N, C_xyz = xyz.shape
if npoint < N:
# fps_idx = torch.arange(npoint).repeat(xyz.shape[0], 1).int().cuda() # DEBUG ONLY
fps_idx = farthest_point_sample_cuda(xyz, npoint) # [B, npoint]
torch.cuda.empty_cache()
new_xyz = index_points_cuda(xyz, fps_idx) #[B, npoint, 3]
new_points = index_points_cuda(points.transpose(1, 2), fps_idx)
else:
new_xyz = xyz
if fps_only:
return new_xyz.transpose(1, 2), new_points.transpose(1, 2), fps_idx
torch.cuda.empty_cache()
_, idx = knn(xyz.contiguous(), new_xyz) # B, npoint, k
idx = idx.int()
torch.cuda.empty_cache()
grouped_xyz = grouping_operation_cuda(
xyz.transpose(1, 2).contiguous(),
idx).permute(0, 2, 3, 1) # [B, npoint, k, C_xyz]
torch.cuda.empty_cache()
grouped_xyz_norm = grouped_xyz - new_xyz.view(B, npoint, 1,
C_xyz) # [B, npoint, k, 3]
grouped_xyz_norm = grouped_xyz_norm.permute(
0, 3, 1, 2).contiguous() # [B, 3, npoint, k]
torch.cuda.empty_cache()
grouped_points = grouping_operation_cuda(points.contiguous(),
idx) #B, C, npoint, k
if cat_xyz_feature:
new_points = torch.cat([grouped_xyz_norm, grouped_points],
dim=1) # [B, C+C_xyz, npoint, k]
else:
new_points = grouped_points # [B, C+C_xyz, npoint, k]
return new_xyz.transpose(1, 2), grouped_xyz_norm, new_points, idx
def __init__(self, r, k, knn=False):
self.r = r
self.k = k
if knn:
self.use_knn = True
self.knn = KNN(k=k, transpose_mode=True)
else:
self.use_knn = False
def sample_and_group_cuda(npoint, k, xyz, points, cat_xyz_feature=True):
"""
Input:
npoint:
k:
xyz: input points position data, [B, N, 3]
points: input points data, [B, C, N]
Return:
new_xyz: sampled points position data, [B, 3, npoint]
new_points: sampled points data, [B, C+C_xyz, npoint, k]
grouped_xyz_norm: sampled relative points position data, [B, 3, npoint, k]
"""
k = min(npoint, k)
knn = KNN(k=k, transpose_mode=True)
B, N, C_xyz = xyz.shape
if npoint < N:
fps_idx = farthest_point_sample_cuda(xyz, npoint) # [B, npoint]
torch.cuda.empty_cache()
new_xyz = index_points_cuda(xyz, fps_idx) #[B, npoint, 3]
else:
new_xyz = xyz
torch.cuda.empty_cache()
_, idx = knn(xyz.contiguous(), new_xyz) # B, npoint, k
idx = idx.int()
torch.cuda.empty_cache()
grouped_xyz = grouping_operation_cuda(
xyz.transpose(1, 2).contiguous(),
idx).permute(0, 2, 3, 1) # [B, npoint, k, C_xyz]
torch.cuda.empty_cache()
try:
# DEBUG: when using the mixed-trans, some last voxels may have less points
grouped_xyz_norm = grouped_xyz - new_xyz.view(-1, min(
npoint, N), 1, C_xyz) # [B, npoint, k, 3]
except:
import ipdb
ipdb.set_trace()
grouped_xyz_norm = grouped_xyz_norm.permute(
0, 3, 1, 2).contiguous() # [B, 3, npoint, k]
torch.cuda.empty_cache()
grouped_points = grouping_operation_cuda(points.contiguous(),
idx) #B, C, npoint, k
if cat_xyz_feature:
new_points = torch.cat([grouped_xyz_norm, grouped_points],
dim=1) # [B, C+C_xyz, npoint, k]
else:
new_points = grouped_points # [B, C+C_xyz, npoint, k]
return new_xyz.transpose(1, 2), grouped_xyz_norm, new_points
def stem_knn(xyz, points, k):
knn = KNN(k=k, transpose_mode=True)
xyz = xyz.permute([0, 2, 1])
_, idx = knn(xyz.contiguous(),
xyz) # xyz: [bs, npoints, coord] idx: [bs, npoint, k]
idx = idx.int()
# take in [B, 3, N]
grouped_xyz = grouping_operation_cuda(
xyz.transpose(1, 2).contiguous(), idx) # [bs, xyz, n_point, k]
grouped_points = grouping_operation_cuda(points.contiguous(),
idx) #B, C, npoint, k)
return grouped_xyz, grouped_points
def sample_and_group_cuda(npoint, k, xyz, points):
"""
Input:
npoint:
k:
xyz: input points position data, [B, N, 3]
points: input points data, [B, C, N]
Return:
new_xyz: sampled points position data, [B, 3, npoint]
new_points: sampled points data, [B, C+C_xyz, npoint, k]
grouped_xyz_norm: sampled relative points position data, [B, 3, npoint, k]
"""
k = min(npoint, k)
knn = KNN(k=k, transpose_mode=True)
B, N, C_xyz = xyz.shape
if npoint < N:
fps_idx = farthest_point_sample_cuda(xyz, npoint) # [B, npoint]
torch.cuda.empty_cache()
new_xyz = index_points_cuda(xyz, fps_idx) #[B, npoint, 3]
else:
new_xyz = xyz
torch.cuda.empty_cache()
_, idx = knn(xyz.contiguous(), new_xyz) # B, npoint, k
idx = idx.int()
torch.cuda.empty_cache()
grouped_xyz = grouping_operation_cuda(
xyz.transpose(1, 2).contiguous(),
idx).permute(0, 2, 3, 1) # [B, npoint, k, C_xyz]
#print(grouped_xyz.size())
torch.cuda.empty_cache()
grouped_xyz_norm = grouped_xyz - new_xyz.view(B, npoint, 1,
C_xyz) # [B, npoint, k, 3]
grouped_xyz_norm = grouped_xyz_norm.permute(
0, 3, 1, 2).contiguous() # [B, 3, npoint, k]
torch.cuda.empty_cache()
grouped_points = grouping_operation_cuda(points.contiguous(),
idx) #B, C, npoint, k
new_points = torch.cat([grouped_xyz_norm, grouped_points],
dim=1) # [B, C+C_xyz, npoint, k]
return new_xyz.transpose(1, 2), grouped_xyz_norm, new_points
def forward(self, input_p, input_x, instance=None, instance_relation=None):
'''
input_p: B, 3, npoint
input_x: B, in_dim, npoint
'''
'''
how to use the instance information:
1. use it as guidance of the attention, mask the knns points with different instance label
2. directly random choose points of same instance label as attention receptive field
3. attend to the instance center
'''
INSTANCE_SCHEME = 3
B, in_dim, npoint = list(input_x.size())
n_sample = self.n_sample
k = min(n_sample, npoint)
h = self.nhead
res = input_x
input_p = input_p.permute([0, 2, 1])
ori_input_p = input_p
if instance is not None and INSTANCE_SCHEME == 1:
# knn more points for sampling
knn_sample_more_ratio = 2
enlarged_k = k * knn_sample_more_ratio
self.knn = KNN(k=enlarged_k, transpose_mode=True)
else:
self.knn = KNN(k=k, transpose_mode=True)
# DEBUG: error here is that in the last block only 4-points;
# however the knn still gives 16 idxs
# so when n-point is smaller than the k(n_smaple)
# if npoint < self.n_sample:
# self.knn = KNN(k=npoint, transpose_mode=True)
# else:
# self.knn = KNN(k=n_sample, transpose_mode=True)
# pass # regular case
# DEBUG ONLY: using the input_x: feature space knn!
# _, idx = self.knn(input_x.transpose(1,2), input_x.transpose(1,2))
if instance is not None:
if INSTANCE_SCHEME == 3:
'''
Ver3.0: use cur instance center as knn center,
calc the instance center, and weighting the cur-idx and the instance center idx
ERROR:
- all points of the same instance will have the same idxes? and all are cloest N points to centroid
- if use weiighted center and coord, However, need to do N-pointx KNN, will be slow...
'''
ori_input_p = input_p.clone()
# where = torch.where(instance[0] == 1)
# instance_xyz = input_p[:,where,:].mean(dim=1) # get [bs, 3] centroid for cur instance
for i_bs in range(instance.shape[0]):
for v in torch.unique(instance[i_bs]):
tmp_idx = torch.where(instance[i_bs] == v)[0]
ins_center = input_p[:, tmp_idx, :].mean(
dim=1) # the centroids for each intsance
# average cur point and the instance center
alpha = 0.999
input_p[:,
tmp_idx, :] = alpha * input_p[:, tmp_idx, :] + (
1 - alpha) * ins_center.unsqueeze(
1) # [bs, n_cur_ins, 3] + [bs, 1, 3]
_, idx = self.knn(ori_input_p.contiguous(), ori_input_p)
_, idx2 = self.knn(ori_input_p.contiguous(), input_p)
print((idx == idx2).int().sum() / idx.nelement())
else:
_, idx = self.knn(input_p.contiguous(), input_p)
else:
_, idx = self.knn(input_p.contiguous(), input_p)
idx = idx.int()
if INSTANCE_SCHEME == 1:
'''
Ver1.0(Naive Version): mask the knn(instance label as auxiliary filter)
older version of the instance mask
directly ck if knn grouped point within the same pointset
then mask if not in
'''
if instance is not None:
# print('start processing the instance mask')
masks = []
for i_bs, idx_cur_bs in enumerate(idx):
# [4096, 16] cur_bs_idx
# [4096]: instance_label[i_bs]
mask = instance[i_bs][idx_cur_bs.long()] # [4096, 2*k]
mask = mask - mask[:, 0].unsqueeze(
-1
) # get the 1st column(the 1st element in k-elements is itself)
mask = (mask == 0).int() # acuiqre the 0-1 mask
masks.append(mask)
masks = torch.stack(masks)
print("mask ratio {:.4f}".format(
masks.sum() / masks.nelement())) # >0.5 means ok
'''
generate bigger knn-idx and mask, then choose the 1st n_sample(16) elements
random sample other points from the latter, and use mask to fill into the 0 ones
get the 1st k idxes that is not 0 in mask
since the mask values are all 0-1, use argsort will return a vector
however, we want smaller idxes in the front
so we give 0 elments a large value to make it appears at last
if use descend=True, biggest idx with 1 will come first
'''
inv_masks = (masks == 0).int()
tmp_inds = torch.arange(masks.shape[2]).repeat(
masks.shape[0], masks.shape[1],
1).to(idx.device) # generate the [1,2,...,enlarged_k] inds
tmp_inds = tmp_inds * masks
tmp_inds = tmp_inds + (
masks.shape[2] + 1
) * inv_masks # fill the places of 0 with value bigger than the maximum value
tmp_inds = torch.argsort(
tmp_inds
)[:, :, :
k] # after argsort, the former elements should be non-zero value with smaller idx
idx = torch.gather(idx, -1, tmp_inds)
idx = idx.int()
# TODO: if nk still does not contain enough elements, the argsort will contain the closet knn result while not instance
elif INSTANCE_SCHEME == 2:
'''
# Ver2.0: directly use the points of the same instance label as neighbors
# random sample k points in the same instance
'''
if instance is not None:
instance_relations = []
for i_bs in range(instance.shape[0]):
instance_inds = [
torch.where(instance[i_bs] == v)[0]
for v in torch.unique(instance[i_bs])
] # torch.where returns a tuple, so use [0] to getthe tensor
instance_relation = torch.full([instance[0].shape[0], k],
-1).to(instance.device)
for i, ins_id in enumerate(instance_inds):
# TODO; stupid pytorch has no func like random.choice
if len(ins_id
) <= 5: # for small outlier points, skip em
continue
try:
perms = torch.multinomial(
ins_id.repeat(len(ins_id), 1).float(),
num_samples=min(k, len(ins_id)),
replacement=False)
except RuntimeError:
import ipdb
ipdb.set_trace()
choices = ins_id[perms]
instance_relation[
instance_inds[i], :choices.shape[1]] = choices
instance_relation[:, 0] = torch.arange(
instance_relation.shape[0])
instance_relations.append(instance_relation)
instance_relations = torch.stack(instance_relations)
# print('replacing the instance_relation')
instance_relation_nonzero_mask = (instance_relations >=
0).int()
instance_relation_zero_mask = (instance_relations < 0).int()
idx = idx * instance_relation_zero_mask + instance_relations * instance_relation_nonzero_mask
idx = idx.int()
# ===================== Deprecated Methods ===========================1
'''
# Ver 2.3: failed version of receiving a instance_relation,
# however, point downsample could not be handled
# the instance feed in here is of the same size as the idxes
# if the num within the same instance group as less than k
# then the instance_relation will contain -1, we will replace these -1s
# with the original idx acquired by knn
if instance_relation is not None:
print('replacing the instance_relation')
# import ipdb; ipdb.set_trace()
instance_relation = instance_relation[:,:,:k]
instance_relation_nonzero_mask = (instance_relation>=0).int()
instance_relation_zero_mask = (instance_relation<0).int()
# idx = idx*instance_relation_zero_mask + instance_relation*instance_relation_nonzero_mask
idx = instance_relation.int()
'''
'''
Ver 2.2: Hash Table-based
1st pack the instance into dict(hash table)
then ref the points within the same scope to replace the knn points
if ont enough points of the same insatcne, keep the knn idxs
'''
'''
# pack the instacne into dict for further reference
if instance is not None:
print('start creating instance dicts')
instance_dicts = []
for i_bs, instance_cur_bs in enumerate(instance):
instance_dict = {}
for ins_idx, ins in enumerate(instance_cur_bs):
if ins.item() in instance_dict.keys():
instance_dict[ins.item()].append(ins_idx)
else:
instance_dict[ins.item()] = [ins_idx]
for ins_k in instance_dict.keys():
instance_dict[ins_k] = torch.tensor(instance_dict[ins_k]).to(instance.device)
instance_dicts.append(instance_dict)
l1 = []
for i_bs in range(instance.shape[0]):
l0 = []
for i_point in range(instance.shape[1]):
tmp = torch.zeros([k])
instance_gathered = instance_dicts[i_bs][instance[i_bs][i_point].item()][:k]
tmp[:len(instance_gathered)] = instance_gathered
# idx[i_bs][i_point][:len(instance_gathered)] = instance_gathered
l0.append(tmp)
tmp1 = torch.stack(l0)
l1.append(tmp1)
new_idx = torch.stack(l1)
'''
'''
Ver: 2.1: Naive Version of for-loop replacement
# Too slow version, needs improving
# 1st use knn then use mask the value belongs not to the same instance
instance_masks = []
for i_batch, single_batch_instance in enumerate(instance):
# single_batch_instance: [npoint]
masks_cur_batch = []
for i_point, gathered_points in enumerate(idx[i_batch]):
# gathered_points: [k]
points_with_same_instance = torch.where(single_batch_instance == single_batch_instance[i_point])[0]
# ck if the grouped idxes are within the same idxes
cur_mask = torch.tensor([g.item() in points_with_same_instance for g in gathered_points])
masks_cur_batch.append(cur_mask)
masks_cur_batch = torch.stack(masks_cur_batch)
instance_masks.append(masks_cur_batch)
instance_masks = torch.stack(instance_masks)
'''
# ==========================================================================================
grouped_input_p = grouping_operation_cuda(
input_p.transpose(1, 2).contiguous(), idx) # [bs, xyz, npoint, k]
if self.pre_ln:
input_x = self.ln_top(input_x.transpose(1, 2)).transpose(1, 2)
input_x = self.linear_top(input_x)
# TODO: apply the layer-norm
# however the original is [bs, dim, npoint]
if self.pre_ln:
input_x = self.ln_attn(input_x.transpose(1, 2)).transpose(1, 2)
# grouped_input_x = index_points(input_x.permute([0,2,1]), idx.long()).permute([0,3,1,2])
# grouped_input_x = grouping_operation_cuda(input_x.contiguous(), idx) # [bs, xyz, npoint, K]
phi = self.phi(input_x)
phi = phi[:, :, :, None].repeat(1, 1, 1, k)
psi = grouping_operation_cuda(self.psi(input_x).contiguous(), idx)
alpha = grouping_operation_cuda(self.alpha(input_x).contiguous(),
idx) # [bs, xyz, npoint, k]
relative_xyz = input_p.permute([0, 2, 1])[:, :, :,
None] - grouped_input_p
pos_encoding = self.delta(relative_xyz) # [bs, dims, npoint, k]
if self.use_vector_attn:
# the attn_map: [vector_dim];
# the alpha: [out_dim]
attn_map = F.softmax(self.gamma(phi - psi + pos_encoding),
dim=-1) # [B, Dim, N, k]
# if instance is not None: # apply mask
# attn_map = attn_map*(masks.unsqueeze(1))
y = attn_map.repeat(1, self.out_dim // self.vector_dim, 1,
1) * (alpha + pos_encoding)
y = y.sum(dim=-1)
else:
phi = phi.reshape(B, h, self.out_dim // h, npoint, k)
psi = psi.reshape(B, h, self.out_dim // h, npoint, k)
attn_map = F.softmax(
(phi * psi).reshape(B, self.out_dim, npoint, k) + pos_encoding,
dim=-1)
y = attn_map * (alpha + pos_encoding)
y = y.sum(dim=-1)
if self.pre_ln:
y = self.ln_down(y.transpose(1, 2)).transpose(1, 2)
y = self.linear_down(y)
return y + res, attn_map.detach().cpu().data
from dataset import *
import my_utils
my_utils.plant_seeds(randomized_seed=False)
from sklearn.neighbors import NearestNeighbors
from ply import *
import os
os.environ['CUDA_VISIBLE_DEVICES'] = "0"
from model_err_nfeaPointnet2 import *
import json
import datetime
from LaplacianLoss import *
from knn_cuda import KNN
knn = KNN(k=1, transpose_mode=True)
import time
import warnings
warnings.filterwarnings("ignore")
# =============PARAMETERS======================================== #
lambda_laplace = 0.005
lambda_ratio = 0.005
parser = argparse.ArgumentParser()
parser.add_argument('--batchSize', type=int, default=32, help='input batch size')
parser.add_argument('--workers', type=int, help='number of data loading workers', default=48)
parser.add_argument('--nepoch', type=int, default=75, help='number of epochs to train for')
parser.add_argument('--model', type=str, default='', help='optional reload model path')
parser.add_argument('--env', type=str, default="3DCODED_unsupervised", help='visdom environment')
parser.add_argument('--laplace', type=int, default=1, help='regularize towords 0 curvature, or template curvature')
def forward(self, input):
r"""
Forward pass
Parameters
----------
input: torch.Tensor, shape (B, N, d_in)
input points
Returns
-------
torch.Tensor, shape (B, num_classes, N)
segmentation scores for each point
"""
N = input.size(1)
d = self.decimation
coords = input[..., :3].clone().cpu()
x = self.fc_start(input).transpose(-2, -1).unsqueeze(-1)
x = self.bn_start(x) # shape (B, d, N, 1)
decimation_ratio = 1
# <<<<<<<<<< ENCODER
x_stack = []
permutation = torch.randperm(N)
coords = coords[:, permutation]
x = x[:, :, permutation]
for lfa in self.encoder:
# at iteration i, x.shape = (B, N//(d**i), d_in)
startenc = time.time()
x = lfa(coords[:, :N // decimation_ratio], x)
x_stack.append(x.clone())
decimation_ratio *= d
x = x[:, :, :N // decimation_ratio]
endenc = time.time()
#print("each dec time is",endenc-startenc)
# # >>>>>>>>>> ENCODER
x = self.mlp(x)
#print("encoder time is",end-start)
# <<<<<<<<<< DECODER
for mlp in self.decoder:
knn = KNN(k=1, transpose_mode=True)
_, neighbors = knn(
coords[:, :N //
decimation_ratio].cuda().contiguous(), # original set
coords[:, :d * N //
decimation_ratio].cuda().contiguous() # upsampled set
) # shape (B, N, 1)
neighbors = neighbors.to(self.device)
extended_neighbors = neighbors.unsqueeze(1).expand(
-1, x.size(1), -1, 1)
x_neighbors = torch.gather(x, -2, extended_neighbors)
x = torch.cat((x_neighbors, x_stack.pop()), dim=1)
x = mlp(x)
decimation_ratio //= d
# >>>>>>>>>> DECODER
# inverse permutation
x = x[:, :, torch.argsort(permutation)]
scores = self.fc_end(x)
return scores.squeeze(-1)
def __init__(self,k=16):
super(get_edge_feature,self).__init__()
self.KNN=KNN(k=k+1,transpose_mode=False)
self.k=k
def create_knn_adj_mat(features,
k,
weighted=False,
n_jobs=None,
algorithm='auto',
threshold=None,
use_gpu=False):
"""
Create a directed normalized adjacency matrix from input nodes based on k-nearest neighbours
Parameters:
features (numpy array of node features): array of size N*M (N nodes, M feature size)
k (int): number of neighbours to find
weighted (bool): set to True for weighted adjacency matrix (based on Euclidean distance)
n_jobs (int): number of jobs to deploy if GPU is not used
algorithm (str): Choose between auto, ball_tree, kd_tree or brute
threshold (float): Cutoff value for the Euclidean distance
use_gpu (bool): Indicates whether GPU is to be used for the KNN algorithm
Returns:
(coo matrix): adjacency matrix as a sparse coo matrix
"""
t_start = datetime.now()
if use_gpu:
features_extra_dim = np.expand_dims(features, axis=2)
knn = KNN(k=k, transpose_mode=True)
# Find the k nearest neighbours and their distance
dist, idx = knn(
torch.from_numpy(features_extra_dim).cuda(),
torch.from_numpy(features_extra_dim).clone().cuda())
torch.cuda.empty_cache()
del features_extra_dim
idx = idx.cpu()
dist = dist.cpu()
# Clean up the indices and distances
dist = dist.flatten()
# Create tuples of indices where an edge exists
rows = np.repeat(np.arange(features.shape[0]), k)
columns = idx.flatten()
non_zero_indices = tuple(np.stack((rows, columns)))
del rows
del columns
del idx
# Remove edges where the distance is higher than the threshold
if threshold:
indices_to_remove = dist > threshold
indices_to_remove = np.where(indices_to_remove)
non_zero_indices = tuple(
np.delete(non_zero_indices, indices_to_remove, 1))
dist = np.delete(dist, indices_to_remove[0], 0)
del indices_to_remove
if weighted:
# Create zero matrix as the initial adjacency matrix
adj_mat_weighted = np.zeros((features.shape[0], features.shape[0]),
dtype=np.float32)
# Fill in the adjacency matrix with node distances
adj_mat_weighted[non_zero_indices] = dist
non_zero_indices = np.nonzero(adj_mat_weighted)
# Take reciprocal of non-zero elements to associate lower weight to higher distances
adj_mat_weighted[
non_zero_indices] = 1 / adj_mat_weighted[non_zero_indices]
# Normalize rows
coo_matrix = sp.coo_matrix(adj_mat_weighted)
normalized_coo_matrix = normalize(coo_matrix)
# DGL requires self loops
normalized_coo_matrix = normalized_coo_matrix + sp.eye(
normalized_coo_matrix.shape[0])
t_end = datetime.now()
logger.debug("it took {} to create the graph".format(t_end -
t_start))
return normalized_coo_matrix
else:
# Create eye matrix as the initial adjacency matrix
adj_mat_binary = np.zeros((features.shape[0], features.shape[0]))
# Create the binary adjacency matrix
adj_mat_binary[non_zero_indices] = 1
t_end = datetime.now()
logger.debug("it took {} to create the graph".format(t_end -
t_start))
return sp.coo_matrix(adj_mat_binary)
else:
# initialize and fit nearest neighbour algorithm
neigh = NearestNeighbors(n_neighbors=k,
n_jobs=n_jobs,
algorithm=algorithm)
neigh.fit(features)
# Obtain matrix with distance of k-nearest points
adj_mat_weighted = np.array(
sp.coo_matrix(neigh.kneighbors_graph(features, k,
mode='distance')).toarray())
if threshold:
indices_to_zero = adj_mat_weighted > threshold
adj_mat_weighted[indices_to_zero] = 0
non_zero_indices = np.nonzero(adj_mat_weighted)
if weighted:
# Take reciprocal of non-zero elements to associate lower weight to higher distances
adj_mat_weighted[
non_zero_indices] = 1 / adj_mat_weighted[non_zero_indices]
# Normalize rows
adj_mat_weighted = sp.coo_matrix(adj_mat_weighted)
normalized_coo_matrix = normalize(adj_mat_weighted)
# DGL requires self loops
normalized_coo_matrix = normalized_coo_matrix + sp.eye(
normalized_coo_matrix.shape[0])
del adj_mat_weighted
t_end = datetime.now()
logger.debug("it took {} to create the graph".format(t_end -
t_start))
return normalized_coo_matrix
# Obtain the binary adjacency matrix
adj_mat_binary = adj_mat_weighted
adj_mat_binary[non_zero_indices] = 1
adj_mat_binary = adj_mat_binary + np.eye(adj_mat_binary.shape[0])
del adj_mat_weighted
t_end = datetime.now()
logger.debug("it took {} to create the graph".format(t_end - t_start))
return sp.coo_matrix(adj_mat_binary)
def knn(ref, query, n, k) -> torch.Tensor:
ref = ref.view(1,n,3)
query = query.view(1,n,3)
d, I = KNN(k, transpose_mode=True)(ref=ref, query=query)
return d.view(n,k), I.view(n*k)
def knn_grad(ref, query, n, k) -> torch.Tensor: #NOTE output tensor shape [n,k,3]
ref = ref.view(1,n,3)
query = query.view(1,n,3)
d, I = KNN(ref=ref, query=query)
diff = query.view(n,1,3) - ref[0, I.view(-1),:].view(n,k,3) #shape [n,k,3]
return diff.view(n,k,3), I
def _knn_point(self, nsample, ref, query):
knn = KNN(k=nsample, transpose_mode=True)
dist, indx = knn(ref, query)
return dist, indx
def __init__(self, radius=1.0):
super(Loss, self).__init__()
self.radius = radius
self.knn_uniform = KNN(k=2, transpose_mode=True)
self.knn_repulsion = KNN(k=20, transpose_mode=True)
def __init__(self, sharpen=1.0):
super().__init__()
self.k = 50
self.knn = KNN(self.k, transpose_mode=True)
self.sharpen = sharpen
def forward(self, input_p, input_x):
'''
input_p: B, 3, npoint
input_x: B, in_dim, npoint
'''
B, in_dim, npoint = list(input_x.size())
n_sample = self.n_sample
k = min(n_sample, npoint)
h = self.nhead
res = input_x
input_p = input_p.permute([0, 2, 1])
# DEBUG: error here is that in the last block only 4-points;
# however the knn still gives 16 idxs
if npoint < self.n_sample:
self.knn = KNN(k=npoint, transpose_mode=True)
# DEBUG ONLY: using the input_x: feature space knn!
# _, idx = self.knn(input_x.transpose(1,2), input_x.transpose(1,2))
_, idx = self.knn(input_p.contiguous(), input_p)
idx = idx.int()
grouped_input_p = grouping_operation_cuda(
input_p.transpose(1, 2).contiguous(), idx) # [bs, xyz, npoint, k]
if self.pre_ln:
input_x = self.ln_top(input_x.transpose(1, 2)).transpose(1, 2)
input_x = self.linear_top(input_x)
# TODO: apply the layer-norm
# however the original is [bs, dim, npoint]
if self.pre_ln:
input_x = self.ln_attn(input_x.transpose(1, 2)).transpose(1, 2)
# grouped_input_x = index_points(input_x.permute([0,2,1]), idx.long()).permute([0,3,1,2])
# grouped_input_x = grouping_operation_cuda(input_x.contiguous(), idx) # [bs, xyz, npoint, K]
phi = self.phi(input_x)
phi = phi[:, :, :, None].repeat(1, 1, 1, k)
psi = grouping_operation_cuda(self.psi(input_x).contiguous(), idx)
alpha = grouping_operation_cuda(self.alpha(input_x).contiguous(),
idx) # [bs, xyz, npoint, k]
relative_xyz = input_p.permute([0, 2, 1])[:, :, :,
None] - grouped_input_p
pos_encoding = self.delta(relative_xyz) # [bs, dims, npoint, k]
if self.use_vector_attn:
# the attn_map: [vector_dim];
# the alpha: [out_dim]
attn_map = F.softmax(self.gamma(phi - psi + pos_encoding), dim=-1)
y = attn_map.repeat(1, self.out_dim // self.vector_dim, 1,
1) * (alpha + pos_encoding)
y = y.sum(dim=-1)
else:
phi = phi.reshape(B, h, self.out_dim // h, npoint, k)
psi = psi.reshape(B, h, self.out_dim // h, npoint, k)
attn_map = F.softmax(
(phi * psi).reshape(B, self.out_dim, npoint, k) + pos_encoding,
dim=-1)
y = attn_map * (alpha + pos_encoding)
y = y.sum(dim=-1)
if self.pre_ln:
y = self.ln_down(y.transpose(1, 2)).transpose(1, 2)
y = self.linear_down(y)
return y + res, attn_map.detach().cpu().data
def __init__(self, in_dim, is_firstlayer=False, n_sample=16):
super().__init__()
'''
Point Transformer Layer
in_dim: feature dimension of the input feature x
out_dim: feature dimension of the Point Transformer Layer(currently same with hidden-dim)
[?] - not sure how to set hidden. the paper only gives the out
'''
self.in_dim = in_dim
self.is_firstlayer = is_firstlayer
# TODO: set the hidden/vector/out_dims
self.hidden_dim = in_dim
# self.out_dim = min(4*in_dim, 512)
self.out_dim = in_dim
self.vector_dim = self.out_dim
self.n_sample = n_sample
# whether use BN or LN or None
# 0 - None
# 1 - BN
# 2 - LN
self.use_bn = 1
# use transformer-like preLN before the attn & ff layer
self.pre_ln = False
# whether to use the vector att or the original attention
self.use_vector_attn = True
self.nhead = 4
self.linear_top = nn.Sequential(
nn.Conv1d(in_dim, self.hidden_dim, 1),
# TransposeLayerNorm(self.hidden_dim),
nn.BatchNorm1d(self.hidden_dim) if self.use_bn else nn.Identity())
self.linear_down = nn.Sequential(
nn.Conv1d(self.out_dim, self.in_dim, 1),
# TransposeLayerNorm(self.in_dim),
nn.BatchNorm1d(self.in_dim) if self.use_bn else nn.Identity())
self.phi = nn.Sequential(
nn.Conv1d(self.hidden_dim, self.out_dim, 1),
# nn.BatchNorm1d(self.out_dim) if self.use_bn else nn.Identity()
)
self.psi = nn.Sequential(
nn.Conv1d(self.hidden_dim, self.out_dim, 1),
# nn.BatchNorm1d(self.out_dim) if self.use_bn else nn.Identity()
)
self.alpha = nn.Sequential(
nn.Conv1d(self.hidden_dim, self.out_dim, 1),
# nn.BatchNorm1d(self.out_dim) if self.use_bn else nn.Identity()
)
self.gamma = nn.Sequential(
nn.Conv2d(self.out_dim, self.hidden_dim, 1),
# TransposeLayerNorm(self.hidden_dim),
nn.BatchNorm2d(self.hidden_dim) if self.use_bn else nn.Identity(),
nn.ReLU(),
nn.Conv2d(self.hidden_dim, self.vector_dim, 1),
# TransposeLayerNorm(self.vector_dim),
nn.BatchNorm2d(self.vector_dim) if self.use_bn else nn.Identity())
self.delta = nn.Sequential(
nn.Conv2d(3, self.hidden_dim, 1),
# TransposeLayerNorm(self.hidden_dim),
nn.BatchNorm2d(self.hidden_dim) if self.use_bn else nn.Identity(),
nn.ReLU(),
nn.Conv2d(self.hidden_dim, self.out_dim, 1),
nn.BatchNorm2d(self.out_dim) if self.use_bn else nn.Identity()
# TransposeLayerNorm(self.out_dim),
)
if self.pre_ln:
self.ln_top = nn.LayerNorm(self.in_dim)
self.ln_attn = nn.LayerNorm(self.hidden_dim)
self.ln_down = nn.LayerNorm(self.out_dim)
self.knn = KNN(k=n_sample, transpose_mode=True)
def forward(self, input_p, input_x):
'''
input_p: B, 3, npoint
input_x: B, in_dim, npoint
'''
B, in_dim, npoint = list(input_x.size()) # npoint: the input point-num
n_sample = self.n_sample # the knn-sample num cur block
k = min(n_sample, npoint) # denoting the num_point cur layer
if not self.use_vector_attn:
h = self.nhead # only used in non-vextor attn
input_p = input_p.permute([0, 2, 1]) # [B, npoint, 3]
self.register_buffer('in_xyz_map', input_p)
if self.fps_rate is not None:
npoint = npoint // self.fps_rate
fps_idx = farthest_point_sample_cuda(input_p, npoint)
torch.cuda.empty_cache()
input_p_fps = index_points_cuda(input_p, fps_idx) # [B. M, 3]
if self.SKIP_ALL:
input_p_reduced = input_p_fps.transpose(1, 2)
input_x_reduced = index_points_cuda(
self.tmp_linear(input_x).transpose(1, 2),
fps_idx).transpose(1, 2)
return input_p_reduced, input_x_reduced
else:
input_p_fps = input_p
input_x_fps = input_x
res = input_x # [B, dim, M]
if self.USE_KNN:
self.knn = KNN(k=k, transpose_mode=True)
_, idx = self.knn(input_p.contiguous(), input_p_fps.contiguous())
idx = idx.int() # [bs, npoint, k]
else:
idx = query_ball_point_cuda(
self.radius, k, input_p.contiguous(),
input_p_fps.contiguous()) # [bs, npoint, k]
grouped_input_p = grouping_operation_cuda(
input_p.transpose(1, 2).contiguous(), idx) # [bs, xyz, npoint, k]
grouped_input_x = grouping_operation_cuda(
input_x.contiguous(), idx) # [bs, hidden_dim, npoint, k]
self.register_buffer('neighbor_map', idx)
# TODO: define proper r for em
# query_idx = query_ball_point_cuda(radius, k, coord, coord) # [bs, npoint, k]
# self.knn = KNN(k=k, transpose_mode=True)
# _, knn_idx = self.knn(input_p.contiguous(), input_p)
# import ipdb; ipdb.set_trace()
if self.fps_rate is not None:
if self.SKIP_ATTN:
pass # only apply linear-top for ds blocks
else:
input_x = self.linear_top(input_x)
else:
if self.SKIP_ATTN:
pass # only apply linear-top for ds blocks
else:
input_x = self.linear_top(input_x)
# input_x = self.linear_top(input_x)
if self.SKIP_ATTN:
# import ipdb; ipdb.set_trace()
# out_dim should be the same with in_dim, since here contains no TD
if self.POS_ENCODING:
relative_xyz = input_p_fps.permute(
[0, 2, 1])[:, :, :, None] - grouped_input_p
pos_encoding = self.delta(
relative_xyz) # [bs, dims, npoint, k]
if self.CAT_POS:
alpha = self.alpha(
torch.cat([grouped_input_x, relative_xyz], dim=1))
else: # use sum
alpha = self.alpha(grouped_input_x + pos_encoding)
else:
alpha = self.alpha(grouped_input_x)
# alpha = grouping_operation_cuda(self.alpha(input_x).contiguous(), idx)
y = alpha.max(dim=-1)[0]
# y = alpha.sum(dim=-1)
y = self.linear_down(y)
if self.fps_rate is not None:
input_p_reduced = input_p_fps.transpose(1, 2)
# WRONG!: noneed for applying fps_idx here
# input_x_reduced = index_points_cuda(y.transpose(1,2), fps_idx).transpose(1,2) # [B, dim, M]
input_x_reduced = y
return input_p_reduced, input_x_reduced
else:
input_p_reduced = input_p_fps.transpose(1, 2)
input_x_reduced = y + res
return input_p_reduced, input_x_reduced
# when downsampling the TRBlock
# should use downsampled qkv here, so use input_x_fps
# as for normal block, input_x and input_x_fps are the same
if self.fps_rate is not None:
input_x_fps = index_points_cuda(input_x.transpose(
1, 2), fps_idx).transpose(
1, 2) # it is only used for tr-like downsample block
phi = self.phi(input_x_fps)
else:
phi = self.phi(input_x)
phi = phi[:, :, :, None].repeat(1, 1, 1, k)
psi = grouping_operation_cuda(self.psi(input_x).contiguous(), idx)
self.skip_knn = True
alpha = grouping_operation_cuda(self.alpha(input_x).contiguous(),
idx) # [bs, xyz, npoint, k]
if self.POS_ENCODING:
relative_xyz = input_p_fps.permute(
[0, 2, 1])[:, :, :, None] - grouped_input_p
pos_encoding = self.delta(relative_xyz) # [bs, dims, npoint, k]
if self.use_vector_attn:
# the attn_map: [vector_dim];
# the alpha: [out_dim]
if self.POS_ENCODING:
# if V_POS and QK_POS is both false, then apply all pos_encoding
assert (
self.V_POS_ONLY and self.QK_POS_ONLY
) is False # only one of the V_ONLY and QK_ONLY should be applied
if self.V_POS_ONLY:
attn_map = F.softmax(self.gamma(phi - psi), dim=-1)
else:
attn_map = F.softmax(self.gamma(phi - psi + pos_encoding),
dim=-1)
if self.QK_POS_ONLY:
y = attn_map.repeat(1, self.out_dim // self.vector_dim, 1,
1) * (alpha)
else:
y = attn_map.repeat(1, self.out_dim // self.vector_dim, 1,
1) * (alpha + pos_encoding)
else:
attn_map = F.softmax(self.gamma(phi - psi), dim=-1)
y = attn_map.repeat(1, self.out_dim // self.vector_dim, 1,
1) * (alpha)
if self.MAX_POOL:
y = y.max(dim=-1)[0]
else:
y = y.sum(dim=-1)
else:
assert self.POS_ENCODING == True
phi = phi.reshape(B, h, self.out_dim // h, npoint, k)
psi = psi.reshape(B, h, self.out_dim // h, npoint, k)
attn_map = F.softmax(
(phi * psi).reshape(B, self.out_dim, npoint, k) + pos_encoding,
dim=-1)
y = attn_map * (alpha + pos_encoding)
y = y.sum(dim=-1)
self.register_buffer('attn_map', attn_map.mean(dim=1))
y = self.linear_down(y)
if self.fps_rate is not None:
input_p_reduced = input_p_fps.transpose(1, 2)
# input_x_reduced = index_points_cuda(y.transpose(1,2), fps_idx).transpose(1,2) # [B, dim, M]
input_x_reduced = y
return input_p_reduced, input_x_reduced
else:
input_p_reduced = input_p_fps.transpose(1, 2)
input_x_reduced = y + res
return input_p_reduced, input_x_reduced
def forward(self, x: torch.Tensor):
# x shape should be B x 3 x N
if self.input_shape == "bnc":
x = x.permute(0, 2, 1)
if x.shape[1] != 3:
raise RuntimeError(
"shape of x must be of [Batch x 3 x NumInPoints]")
y = F.relu(self.bn1(self.conv1(x)))
y = F.relu(self.bn2(self.conv2(y)))
y = F.relu(self.bn3(self.conv3(y)))
y = F.relu(self.bn4(self.conv4(y)))
y = F.relu(self.bn5(self.conv5(y))) # Batch x 128 x NumInPoints
# Max pooling for global feature vector:
y = torch.max(y, 2)[0] # Batch x 128
y = F.relu(self.bn_fc1(self.fc1(y)))
y = F.relu(self.bn_fc2(self.fc2(y)))
y = F.relu(self.bn_fc3(self.fc3(y)))
y = self.fc4(y)
y = y.view(-1, 3, self.num_out_points)
# Simplified points
simp = y
match = None
proj = None
# Projected points
if self.training:
if not self.skip_projection:
proj = self.project(point_cloud=x, query_cloud=y)
else:
proj = simp
# Matched points
else: # Inference
# Retrieve nearest neighbor indices
_, idx = KNN(1, transpose_mode=False)(x.contiguous(),
y.contiguous())
"""Notice that we detach the tensors and do computations in numpy,
and then convert back to Tensors.
This should have no effect as the network is in eval() mode
and should require no gradients.
"""
# Convert to numpy arrays in B x N x 3 format. we assume 'bcn' format.
x = x.permute(0, 2, 1).cpu().detach().numpy()
y = y.permute(0, 2, 1).cpu().detach().numpy()
idx = idx.cpu().detach().numpy()
idx = np.squeeze(idx, axis=1)
z = sputils.nn_matching(x,
idx,
self.num_out_points,
complete_fps=self.complete_fps)
# Matched points are in B x N x 3 format.
match = torch.tensor(z, dtype=torch.float32).cuda()
# Change to output shapes
if self.output_shape == "bnc":
simp = simp.permute(0, 2, 1)
if proj is not None:
proj = proj.permute(0, 2, 1)
elif self.output_shape == "bcn" and match is not None:
match = match.permute(0, 2, 1)
match = match.contiguous()
# Assert contiguous tensors
simp = simp.contiguous()
if proj is not None:
proj = proj.contiguous()
if match is not None:
match = match.contiguous()
out = proj if self.training else match
return simp, out
def sample_and_group_cuda(npoint,
k,
xyz,
points,
instance=None,
instance_relation=None):
"""
Input:
npoint: seems 1/4 of N
k:
xyz: input points position data, [B, N, 3]
points: input points data, [B, C, N]
instance: input_instance, [B,N]
Return:
new_xyz: sampled points position data, [B, 3, npoint]
new_points: sampled points data, [B, C+C_xyz, npoint, k]
grouped_xyz_norm: sampled relative points position data, [B, 3, npoint, k]
new_instance, [B, npoint]
"""
k = min(npoint, k)
knn = KNN(k=k, transpose_mode=True)
B, N, C_xyz = xyz.shape
if npoint < N:
fps_idx = farthest_point_sample_cuda(xyz, npoint) # [B, npoint]
torch.cuda.empty_cache()
new_xyz = index_points_cuda(xyz, fps_idx) #[B, npoint, 3]
else:
new_xyz = xyz
# unsqueeze to [B,N,1] then apply indexing
if instance is not None:
new_instance = index_points_cuda(
instance.unsqueeze(-1).float(), fps_idx).squeeze(-1)
else:
pass
torch.cuda.empty_cache()
_, idx = knn(xyz.contiguous(), new_xyz) # B, npoint, k
idx = idx.int()
torch.cuda.empty_cache()
grouped_xyz = grouping_operation_cuda(
xyz.transpose(1, 2).contiguous(),
idx).permute(0, 2, 3, 1) # [B, npoint, k, C_xyz]
#print(grouped_xyz.size())
torch.cuda.empty_cache()
grouped_xyz_norm = grouped_xyz - new_xyz.view(B, npoint, 1,
C_xyz) # [B, npoint, k, 3]
grouped_xyz_norm = grouped_xyz_norm.permute(
0, 3, 1, 2).contiguous() # [B, 3, npoint, k]
torch.cuda.empty_cache()
grouped_points = grouping_operation_cuda(points.contiguous(),
idx) #B, C, npoint, k
new_points = torch.cat([grouped_xyz_norm, grouped_points],
dim=1) # [B, C+C_xyz, npoint, k]
if instance is not None:
return new_xyz.transpose(1,
2), grouped_xyz_norm, new_points, new_instance
else:
return new_xyz.transpose(1, 2), grouped_xyz_norm, new_points
def knn_point(group_size, point_cloud, query_cloud):
knn_obj = KNN(k=group_size, transpose_mode=False)
dist, idx = knn_obj(point_cloud, query_cloud)
return dist, idx
def __init__(self, num):
super(RepulsionLoss, self).__init__()
self.k = num
self.knn_repulsion = KNN(k=self.k, transpose_mode=True)
