如何计算 PyTorch 中点集和线之间的成对距离？

Question

点集A是一个Nx3矩阵，从具有相同大小的Mx3的两个点集B和C我们可以得到它们之间的线BC 。 现在我想计算从A中的每个点到BC中的每条线的距离。 B是Mx3和C是Mx3 ，那么这些线来自具有相应行的点，所以BC是一个Mx3矩阵。 基本方法计算如下：

D = torch.zeros((N, M), dtype=torch.float32)
for i in range(N):
    p = A[i]  # 1x3
    for j in range(M):
        p1 = B[j] # 1x3
        p2 = C[j] # 1x3
        D[i,j] = torch.norm(torch.cross(p1 - p2, p - p1)) / torch.norm(p1 - p2) 

有没有更快的方法来完成这项工作？ 谢谢。

Answer 1

您可以通过这样做删除for循环（它应该以 memory 为代价加速，除非M和N很小）：

diff_B_C = B - C
diff_A_C = A[:, None] - C
norm_lines = torch.norm(diff_B_C, dim=-1)
cross_result = torch.cross(diff_B_C[None, :].expand(N, -1, -1), diff_A_C, dim=-1)
norm_cross = torch.norm(cross_result, dim=-1)
D = norm_cross / norm_lines

当然，你不需要一步一步来。 我只是想清楚变量名。

注意：如果您不向torch.cross提供dim ，它将使用第一个dim=3如果N=3会给出错误的结果（来自文档）：

如果没有给出 dim ，则默认为找到的第一个尺寸为 3 的尺寸。

如果你想知道，你可以在这里查看我为什么选择expand而不是repeat 。