PyTorch中cdist和sum函数使用详解
torch.cdist
是 PyTorch 中用于计算**两个张量之间的成对距离(pairwise distance)**的函数,常用于点云处理、图神经网络、相似性度量等场景。
基本语法
torch.cdist(x1, x2, p=2.0)
参数说明:
参数 | 说明 |
---|---|
x1 | 一个形状为 [B, M, D] 或 [M, D] 的张量,表示一组点。 |
x2 | 一个形状为 [B, N, D] 或 [N, D] 的张量,表示另一组点。 |
p | 距离范数,默认 p=2.0 表示欧几里得距离(L2 范数),也可以设为 1.0 (曼哈顿距离),或其他值。 |
输出
输出是一个张量,形状为:
- 如果
x1.shape = [M, D]
和x2.shape = [N, D]
,则输出形状为[M, N]
; - 每个
(i, j)
位置表示x1[i]
和x2[j]
之间的距离。
示例
1. 简单的 2D 欧几里得距离
import torchx1 = torch.tensor([[0.0, 0.0], [1.0, 0.0]]) # 2个点
x2 = torch.tensor([[0.0, 1.0], [1.0, 1.0]]) # 2个点dist = torch.cdist(x1, x2, p=2)
print(dist)
输出为:
tensor([[1.0000, 1.4142],[1.4142, 1.0000]])
即:
- x1[0] 与 x2[0] 的距离为 1;
- x1[0] 与 x2[1] 的距离为 sqrt(2),等等。
2. 批量形式(3D Tensor)
x1 = torch.rand(2, 5, 3) # batch=2, 每组5个3D点
x2 = torch.rand(2, 4, 3) # batch=2, 每组4个3D点out = torch.cdist(x1, x2) # 输出形状为 [2, 5, 4]
3. 使用不同范数
torch.cdist(x1, x2, p=1) # 曼哈顿距离
torch.cdist(x1, x2, p=2) # 欧几里得距离(默认)
torch.cdist(x1, x2, p=inf) # 最大维度差
注意事项
x1
和x2
的最后一维(特征维度)必须相同。p=2
时效率最高,其他范数可能会慢一些。- 如果两个张量都很大,这个操作可能非常耗显存。
应用场景举例
- 点云之间距离计算(如 ISS、FPFH、ICP)
- 匹配点对的距离图构建
- KNN 查询
- 图构造(邻接矩阵、相容性矩阵)
torch.sum
是 PyTorch 中用于对张量元素进行求和的函数,功能类似于 NumPy 中的 np.sum
,但可以更灵活地选择维度进行操作。
基本用法
torch.sum(input, dim=None, keepdim=False)
参数说明:
input
:要进行求和的张量;dim
(可选):指定在哪个维度上进行求和;keepdim
(可选):布尔值,是否保留被求和的维度(默认不保留)。
示例讲解
示例 1:对所有元素求和
x = torch.tensor([[1, 2], [3, 4]])
torch.sum(x)
# 输出:tensor(10)
示例 2:指定维度求和
x = torch.tensor([[1, 2], [3, 4]])torch.sum(x, dim=0) # 按列求和:1+3, 2+4
# 输出:tensor([4, 6])torch.sum(x, dim=1) # 按行求和:1+2, 3+4
# 输出:tensor([3, 7])
示例 3:保留维度
x = torch.tensor([[1, 2], [3, 4]])torch.sum(x, dim=1, keepdim=True)
# 输出:tensor([[3], [7]])