PointNet++函数square_distance(src, dst)：计算两组点之间的欧式距离（代码详解）

文章目录

- - 一、计算两组点之间的欧式距离
  - 二、举例
  - 三、中间结果输出

一、计算两组点之间的欧式距离

def square_distance(src, dst):
    """
    Calculate Euclid distance between each two points.

    src^T * dst = xn * xm + yn * ym + zn * zm；
    sum(src^2, dim=-1) = xn*xn + yn*yn + zn*zn;
    sum(dst^2, dim=-1) = xm*xm + ym*ym + zm*zm;
    dist = (xn - xm)^2 + (yn - ym)^2 + (zn - zm)^2
         = sum(src**2,dim=-1)+sum(dst**2,dim=-1)-2*src^T*dst

    Input:
        src: source points, [B, N, C]
        dst: target points, [B, M, C]
    Output:
        dist: per-point square distance, [B, N, M]
    """
    B, N, _ = src.shape
    _, M, _ = dst.shape
    dist = -2 * torch.matmul(src, dst.permute(0, 2, 1))
    dist += torch.sum(src ** 2, -1).view(B, N, 1)
    dist += torch.sum(dst ** 2, -1).view(B, 1, M)
    return dist

🍉解释：

B, N, _ = src.shape：获取输入源点和目标点的形状信息，其中 B 表示批量大小，N 表示源点的数量

_, M, _ = dst.shape：M 表示目标点的数量，C 表示每个点的维度

dist = -2 * torch.matmul(src, dst.permute(0, 2, 1))：这一步计算了两组点之间的叉乘积

dst.permute(0, 2, 1)：将目标点张量 dst 的第二维和第三维进行交换，以便进行点积

同理，src为N x C，dst为M x C，需要将M x C转置成C x M才可以进行点积（N x C）·（C x M）
torch.matmul ：计算源点和目标点之间的点积，结果 dist 是一个形状为 [B, N, M] 的张量，表示每对源点和目标点之间的点积

dist += torch.sum(src ** 2, -1).view(B, N, 1)：计算了源点和目标点的平方和，并将其广播到与 dist 相同的形状

torch.sum(src ** 2, -1)：计算张量 src 中每个点的平方和，src ^2 将 src 中的每个元素都平方，然后 torch.sum 函数对最后一个维度（即 -1 所代表的维度）进行求和，最后一个维度被求和消除。
假设 src 张量的形状是 [B, N, D]，其中 B 表示批量大小，N 表示点的数量，D 表示每个点的维度。那么 torch.sum(src ** 2, -1) 的结果形状将是 [B, N]，其中每个元素表示了原张量中相应位置点的平方和
view(B, N, 1)：对张量调整到[B, N, 1]，以便与后续的计算相兼容

dist += torch.sum(dst ** 2, -1).view(B, 1, M)：将这些平方和加到 dist 上，以完成欧氏距离的计算

dist：张量，函数返回每对源点和目标点之间的欧氏距离的平方，形状为 [B, N, M]

计算欧式距离的平方等价于下方等式

$x_{1}-x_{2})^{2}+(y_{1}-y_{2})^{2}+(z_{1}-z_{2})^{2}$ =
$x_{1}^{2}+y_{1}^{2}+z_{1}^{2}+x_{2}^{2}+y_{2}^{2}+z_{2}^{2}-2x_{1}x_{2}-2y_{1}y_{2}-2z_{1}z_{2}$

二、举例

假设有两组点，分别是 src 和 dst：

import torch

def square_distance(src, dst):
    B, N, _ = src.shape
    _, M, _ = dst.shape
    dist = -2 * torch.matmul(src, dst.permute(0, 2, 1))
    dist += torch.sum(src ** 2, -1).view(B, N, 1)
    dist += torch.sum(dst ** 2, -1).view(B, 1, M)
    return dist

# 定义源点和目标点
src = torch.tensor([[[1, 2, 3], [4, 5, 6]]])  # shape: [1, 2, 3]
dst = torch.tensor([[[7, 8, 9], [10, 11, 12], [13, 14, 15]]])  # shape: [1, 3, 3]

dist = square_distance(src, dst)
print(dist)

结果
在这里插入图片描述
例如
$7-1)^{2}+(8-2)^{2}+(9-3)^{2}=108$

三、中间结果输出

对于dist = -2 * torch.matmul(src, dst.permute(0, 2, 1))

对于torch.sum(src ** 2, -1)

tensor([[14, 77]])：这是一个形状为 (1, 2) 的张量，表示一个批次中有两个源点，每个源点有两个坐标分量。具体地，它包含了以下信息：第一个源点的坐标是 (14, 77)。

对于torch.sum(src ** 2, -1).view(B, N, 1)

tensor([[[14], [77]]])：这是一个形状为 (1, 2, 1) 的张量，表示一个批次中有两个目标点，每个目标点有一个坐标分量。具体地，它包含了以下信息：
第一个目标点的坐标是 (14)
第二个目标点的坐标是 (77)

对于dist += torch.sum(src ** 2, -1).view(B, N, 1)

对于torch.sum(dst ** 2, -1)

对于torch.sum(src ** 2, -1).view(B, N, 1)

对于dist += torch.sum(src ** 2, -1).view(B, N, 1)