mirrors
/
mmcv
mirror of https://github.com/open-mmlab/mmcv.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[Fix] Fix the potential NaN bug in calc_square_dist() (#2356)

pull/2358/head
ChaimZhu 2022-10-26 14:06:28 +08:00 committed by GitHub
parent 9709ff3f8c
commit 3c35bab67c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
mmcv/ops/points_sampler.py
View File

@ -24,16 +24,11 @@ def calc_square_dist(point_feat_a: Tensor,
torch.Tensor: (B, N, M) Square distance between each point pair. torch.Tensor: (B, N, M) Square distance between each point pair.
""" """
num_channel = point_feat_a.shape[-1] num_channel = point_feat_a.shape[-1]
# [bs, n, 1] dist = torch.cdist(point_feat_a, point_feat_b)
a_square = torch.sum(point_feat_a.unsqueeze(dim=2).pow(2), dim=-1)
# [bs, 1, m]
b_square = torch.sum(point_feat_b.unsqueeze(dim=1).pow(2), dim=-1)
corr_matrix = torch.matmul(point_feat_a, point_feat_b.transpose(1, 2))
dist = a_square + b_square - 2 * corr_matrix
if norm: if norm:
dist = torch.sqrt(dist) / num_channel dist = dist / num_channel
else:
dist = torch.square(dist)
return dist return dist