Pointnet++代码详解：square_distance函数 & query_ball_point函数
Pointnet++代码详解
-
- Square_distance函数
- Query_ball_point
Square_distance函数
square_distance函数主要用来在ball query过程中确定每一个点距离采样点的距离。函数输入是两组点，N为第一组点的个数，M为第二组点的个数，C为输入点的通道数（如果是xyz时C=3），返回的是两组点之间两两的欧几里德距离，即N×M的矩阵。由于在训练中数据通常是以Mini-Batch的形式输入的，所以有一个Batch数量的维度为B。
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, _ = src.shape
_, M, _ = dst.shape
dist = -2 * torch.matmul(src, dst.permute(0, 2, 1))
dist += torch.sum(src**2, dim=-1).view(B, N, 1)
dist += torch.sum(dst**2, dim=-1).view(B, 1, M)
return dist
1、torch中的几种乘法
# PyTorch 中的矩阵、向量、标量之间的乘法
# 一、torch.mul()
# 注意：torch.mul() 是支持广播操作
# torch.mul(input, value, out=None)
# 用标量值 value 乘以输入 input 的每个元素，并返回一个新的结果张量。 out = tensor ∗ value
# 如果输入是FloatTensor or DoubleTensor类型，则 value 必须为实数，否则须为整数。【译注：似乎并非如此，无关输入类型，value取整数、实数皆可。】
# 参数：
# input (Tensor) – 输入张量
# value (Number) – 乘到每个元素的数
# out (Tensor, optional) – 输出张量
import torch
a = torch.randn(3)
print("a : ", a) # tensor([-1.6289, 0.2446, -0.3691])
print("a.size() : ", a.size()) # torch.Size([3])
mul_a_100 = torch.mul(a, 100)
print("mul_a_100 : ", mul_a_100) # tensor([-162.8945, 24.4566, -36.9136])
print("mul_a_100.size() : ", mul_a_100.size()) # torch.Size([3])
print("*" * 50)
# 两个张量 input, other 按元素进行相乘，并返回到输出张量。即计算 outi = inputi ∗ otheri
# 两计算张量形状不须匹配，但总元素数须一致。 注意：当形状不匹配时，input的形状作为输入张量的形状。
#
# 参数：
#
# input (Tensor) – 第一个相乘张量
# other (Tensor) – 第二个相乘张量
# out (Tensor, optional) – 结果张量
c = torch.randn(4, 4)
print("c.size() : ", c.size()) # torch.Size([4, 4])
# d = torch.randn(2, 8) # torch.Size([2, 8]) 该形状不符合广播条件
# RuntimeError: The size of tensor a (4) must match the size of tensor b (8) at non-singleton dimension 1
d = torch.randn(1, 4) # 该形状符合广播条件
print("d.size() : ", d.size())
mul_c_d = torch.mul(c, d)
print("mul_c_d.size() : ", mul_c_d.size()) # torch.Size([4, 4])
# 二、torch.mm()
# 注意，torch.mm()不支持广播（broadcast）。
# torch.mm(mat1, mat2, out=None) → Tensor
# 对矩阵mat1和mat2进行相乘。 如果mat1 是一个n×m 张量，mat2 是一个 m×p 张量，将会输出一个 n×p 张量out。
print("^" * 50)
mat1 = torch.randn(2, 3)
print(mat1.size()) # torch.Size([2, 3])
# mat2 = torch.randn(1, 3) # 该形状不支持广播
# print(mat2.size()) # torch.Size([1, 3])
# RuntimeError: size mismatch, m1: [2 x 3], m2: [1 x 3] at /pytorch/aten/src/TH/generic/THTensorMath.cpp:752
mat2 = torch.randn(3, 4)
print(mat2.size()) # torch.Size([3, 4])
mm = torch.mm(mat1, mat2)
print(mm.size()) # torch.Size([2, 4])
# 三、torch.mv()
# 注意，torch.mv()不支持广播（broadcast）
# torch.mv(mat, vec, out=None) → Tensor
# 对矩阵mat和向量vec进行相乘。 如果mat 是一个n×m张量，vec 是一个m元 1维张量，将会输出一个n 元 1维张量。
print("-" * 50)
mat = torch.randn(2, 3)
print(mat.size()) # torch.Size([2, 3])
# vec = torch.randn(2)
# RuntimeError: size mismatch, [2 x 3], [2] at /pytorch/aten/src/TH/generic/THTensorMath.cpp:631
vec = torch.randn(3)
print(vec.size()) # torch.Size([3])
mv = torch.mv(mat, vec)
print(mv.size()) # torch.Size([2])
# 四、torch.dot()
# 注意，torch.dot()不支持广播（broadcast）
# torch.dot(tensor1, tensor2) → Tensor
# 计算两个张量的点乘(内乘),两个张量都为1-D 向量
print("=" * 50)
# x = torch.tensor([2, 3, 2]) # 该形状不支持广播
# print(x.size()) # torch.Size([3])
# RuntimeError: inconsistent tensor size, expected tensor [3] and src [2] to have the same number of elements,
# but got 3 and 2 elements respectively
x = torch.tensor([2, 3])
print(x.size()) # torch.Size([2])
y = torch.tensor([4, 1])
print(y.size()) # torch.Size([2])
dot = torch.dot(x, y)
print(dot) # tensor(11)
print(dot.size()) # torch.Size([])
print("~" * 50)
# 五、torch.matmul()
# 注意：torch.matmul() 支持广播
# torch.matmul(input, other, out=None) → Tensor
# 两个张量的矩阵乘积
# 计算结果取决于张量的维度：
# 1）如果两个张量都是 1 维，返回结果为 the dot product (scalar) 【点乘（标量）】
# 2）如果两个张量都是 2 维，返回结果为 the matrix-matrix product (矩阵乘积)
# 3）如果第一个参数是 1 维，第二个参数是 2 维，为了矩阵乘法的目的，在第一维上加 1（达到扩充维度的目的），
# 矩阵计算完成之后，第一维加上的 1 将会被删掉。
# 4）如果第一个参数是 2 维，第二个参数是 1 维，返回结果为 the matrix-vector product (矩阵向量乘积)
# 5）如果两个参数至少是 1 维且至少一个参数为 N 维（其中N> 2），则返回 batched matrix multiply (批处理矩阵乘法)
# 如果第一个参数是 1 维，则在其维数之前添加 1，以实现批量矩阵乘法并在计算之后删除 1。
# 如果第二个参数是 1 维，则在其维数之前添加 1，以实现批量矩阵乘法并在计算之后删除 1。
# 非矩阵（即批处理）尺寸被广播（因此必须是可广播的）。
# 例如，如果 input 的张量是 j×1×n×m ，
# other 的张量是 k×m×p，
# out 的张量将会是 j×k×n×p
# case 1：vector x vector
tensor1 = torch.randn(3)
print(tensor1.size()) # torch.Size([3])
tensor2 = torch.randn(3)
print(tensor2.size()) # torch.Size([3])
matmul_1_2 = torch.matmul(tensor1, tensor2)
print(matmul_1_2) # tensor(0.2001) -- scalar
print(matmul_1_2.size()) # torch.Size([])
# case 4： matrix x vector (该情况下不支持广播，matrix的列数必须要和vector的行数一致才能进行计算)
tensor3 = torch.randn(3, 4)
print(tensor3.size()) # torch.Size([3, 4])
tensor4 = torch.randn(4)
print(tensor4.size()) # torch.Size([4])
matmul_3_4 = torch.matmul(tensor3, tensor4)
print(matmul_3_4) # tensor([ 0.8020, 0.2547, -1.2333])
print(matmul_3_4.size()) # torch.Size([3])
# case 5：batched matrix x broadcasted vector
a = torch.randn(10, 3, 4)
print(a.size()) # torch.Size([10, 3, 4])
b = torch.randn(4)
print(b.size()) # torch.Size([4])
matmul_a_b = torch.matmul(a, b)
print(matmul_a_b.size()) # torch.Size([10, 3])
# case 5：batched matrix x batched matrix
c = torch.randn(10, 3, 4)
print(c.size()) # torch.Size([10, 3, 4])
d = torch.randn(10, 4, 5)
print(d.size()) # torch.Size([10, 4, 5])
matmul_c_d = torch.matmul(c, d)
print(matmul_c_d.size()) # torch.Size([10, 3, 5])
# case 5：batched matrix x broadcasted matrix
m = torch.randn(10, 3, 4)
print(m.size()) # torch.Size([10, 3, 4])
n = torch.randn(4, 5)
print(n.size()) # torch.Size([4, 5])
matmul_m_n = torch.matmul(m, n)
print(matmul_m_n.size()) # torch.Size([10, 3, 5])
2、torch.permute
permute(dims)
将tensor的维度换位。
参数：参数是一系列的整数，代表原来张量的维度。比如三维就有0，1，2这些dimension。
import torch
import numpy as np
a=np.array([[[1,2,3],[4,5,6]]])
unpermuted=torch.tensor(a)
print(unpermuted.size()) # ——> torch.Size([1, 2, 3])
permuted=unpermuted.permute(2,0,1)
print(permuted.size()) # ——> torch.Size([3, 1, 2])
再比如图片img的size比如是（28，28，3）就可以利用img.permute(2,0,1)得到一个size为（3，28，28）的tensor。
利用这个函数permute（1，3，2）可以把Tensor([[[1,2,3],[4,5,6]]]) 转换成
tensor([[[1., 4.],
[2., 5.],
[3., 6.]]])
如果使用view（1,3,2），可以得到:
tensor([[[1., 2.],
[3., 4.],
[5., 6.]]])
这时候就令人有点困惑了，permute和view的区别又是啥呢？
二维的情况
先用二维tensor作为例子，方便理解。
permute作用为调换Tensor的维度，参数为调换的维度。例如对于一个二维Tensor来说，调用tensor.permute(1,0)意为将1轴（列轴）与0轴（行轴）调换，相当于进行转置。
In [20]: a
Out[20]:
tensor([[0, 1, 2],
[3, 4, 5]])
In [21]: a.permute(1,0)
Out[21]:
tensor([[0, 3],
[1, 4],
[2, 5]])
如果使用view(3,2)或reshape(3,2)，得到的tensor并不是转置的效果，而是相当于将原tensor的元素按行取出，然后按行放入到新形状的tensor中。
In [22]: a.reshape(3,2)
Out[22]:
tensor([[0, 1],
[2, 3],
[4, 5]])
In [23]: a.view(3,2)
Out[23]:
tensor([[0, 1],
[2, 3],
[4, 5]])
高维的情况
一般使用permute的情况都是在更高维的情况下使用，例如对于一个图像batch，其形状为[batch, channel, height, width]，我们可以使用tensor.permute(0,3,2,1)得到形状为[batch, width, height, channel]的tensor.
我们构造一个模拟的batch用于演示。
In [25]: a=torch.arange(2*3*2*1).reshape(2,3,2,1)
In [26]: a
Out[26]:
tensor([[[[ 0], # 这是第0张“图片”的第0号通道的2个元素
[ 1]],
[[ 2], # 这是第0张“图片”的第1号通道的2个元素
[ 3]],
[[ 4], # 这是第0张“图片”的第2号通道的2个元素
[ 5]]],
[[[ 6],
[ 7]],
[[ 8],
[ 9]],
[[10],
[11]]]])
a的形状为[2,3,2,1]，这个batch有2张“图片”，每张图片有3个通道，每个通道为2x1，例如第0张图片的第0号通道为
[[0], [1]]
In [27]: a.permute(0,3,2,1)
Out[27]:
tensor([[[[ 0, 2, 4],
[ 1, 3, 5]]],
[[[ 6, 8, 10],
[ 7, 9, 11]]]])
In [28]: a.permute(0,3,2,1).shape
Out[28]: torch.Size([2, 1, 2, 3])
形状为[2,3,2,1]的batch执行permute(0,3,2,1)交换维度之后，得到的是[2,1,2,3]，即[batch, width, height, channel]
可以理解为，对于一个高维的Tensor执行permute，我们没有改变数据的相对位置，而只是旋转了一下这个(超)立方体。或者也可以说，改变了我们对这个(超)立方体的“观察角度”而已。
3、Pytorch的广播机制：
dist += torch.sum(src**2, dim=-1).view(B, N, 1)
dist += torch.sum(dst**2, dim=-1).view(B, 1, M)
代码当中这两行采用了广播机制。
pytorch中的广播机制和numpy中的广播机制一样, 因为都是数组的广播机制。
ndim表示的是A是几维矩阵，即A.shape=(1,9,4)代表A.ndim=3。
#A.ndim>B.ndim
import numpy as np
# a.shape=(2,2,3,4)
a = np.arange(1,25).reshape((2,3,4))
# b.shape=(3,4)
b = np.arange(1,13).reshape((3,4))
# numpy会将b.shape调整至(2,3,4)
res = a + b
print('===================================')
print(a)
print(a.shape)
print('===================================')
print(b)
print(b.shape)
print('===================================')
print(res)
print(res.shape)
#结果：
===================================
[[[ 1 2 3 4]
[ 5 6 7 8]
[ 9 10 11 12]]
[[13 14 15 16]
[17 18 19 20]
[21 22 23 24]]]
(2, 3, 4)
===================================
[[ 1 2 3 4]
[ 5 6 7 8]
[ 9 10 11 12]]
(3, 4)
===================================
[[[ 2 4 6 8]
[10 12 14 16]
[18 20 22 24]]
[[14 16 18 20]
[22 24 26 28]
[30 32 34 36]]]
(2, 3, 4)
#A.ndim==B.ndim
# a.shape=(1,3,4)
a = np.arange(1,13).reshape((1,3,4))
# b.shape=(2,1,4)
b = np.arange(1,9).reshape((2,1,4))
res = a + b
print('===================================')
# print(a)
print(a)
print('===================================')
# print(b)
print(b)
print('===================================')
# print(res)
print(res)
print('===================================')
#结果：
===================================
[[[ 1 2 3 4]
[ 5 6 7 8]
[ 9 10 11 12]]]
===================================
[[[1 2 3 4]]
[[5 6 7 8]]]
===================================
[[[ 2 4 6 8]
[ 6 8 10 12]
[10 12 14 16]]
[[ 6 8 10 12]
[10 12 14 16]
[14 16 18 20]]]
===================================
Query_ball_point
query_ball_point函数对应于Grouping layer, 这一层使用Ball query方法生成N’个局部区域，根据论文中的意思，这里有两个变量 ，一个是每个区域中点的数量K，另一个是球的半径。这里半径应该是占主导的，会在某个半径的球内找点，上限是K。球的半径和每个区域中点的数量都是人指定的。
query_ball_point函数用于寻找球形领域中的点。输入中radius为球形领域的半径，nsample为每个领域中要采样的点，new_xyz为S个球形领域的中心（由最远点采样在前面得出），xyz为所有的点云；输出为每个样本的每个球形领域的nsample个采样点集的索引[B,S,nsample]。
def query_ball_point(radius, nsample, xyz, new_xyz):
"""
Input:
radius: local region radius
nsample: max sample number in local region
xyz: all points, [B, N, 3]
new_xyz: query points, [B, S, 3] ,s denotes the number of center points
Return:
group_idx: grouped points index, [B, S, nsample]
"""
device = xyz.device
B, N, C = xyz.shape
_, S, _ = new_xyz.shape
group_idx = torch.arange(N, dtype=torch.long).to(device)\
.view(1, 1, N).repeat([B, S, 1])
# sqrdists: [B, S, N] 记录中心点与所有点之间的欧几里德距离
sqrdists = square_distance(new_xyz, xyz)
# 找到所有距离大于radius^2的，其group_idx直接置为N；其余的保留原来的值
group_idx[sqrdists > radius **2] = N
# 做升序排列，前面大于radius^2的都是N，会是最大值，所以会直接在剩下的点中取出前nsample个点
group_idx = group_idx.sort(dim=-1)[0][:, :, :nsample]
# 考虑到有可能前nsample个点中也有被赋值为N的点（即球形区域内不足nsample个点），这种点需要舍弃，直接用第一个点来代替即可
# group_first: [B, S, nsample]， 实际就是把group_idx中的第一个点的值复制到[B, S, nsample]的维度，便利于后面的替换
# 这里要用view是因为group_idx[:, :, 0]取出之后的tensor相当于二维Tensor，因此需要用view变成三维tensor
group_first = group_idx[:, :, 0].view(B, S, 1).repeat([1, 1, nsample])
# 找到group_idx中值等于N的点，会输出0,1构成的三维Tensor，维度为[B,S,nsample]
mask = group_idx == N
# 将这些点的值替换为第一个点的值
group_idx[mask] = group_first[mask]
return group_idx
1、对于group_idx的理解：
group_idx = torch.arange(N, dtype=torch.long).to(device)\
.view(1, 1, N).repeat([B, S, 1])
N指的是一个样本中总的数据点的个数，用torch.arange(N)可以生成tensor([0,1,…,N-1]), 用.to(device)意思是说将生成的tensor([0,1,…,N-1])复制到的xyz所在的设备上，再用.view(1,1,N)则将tesor表示成tesnor([[[0,1,…,N-1]]])即有N列的意思,再用.repeat([B,S,1])则是说将原来的tensor在维度0上复制B个（原先只有1个），在维度1上复制S个，可以理解有B个batch,每个样本有S行N列，所以最后group_idx的维度为[B,S,N], 用代码来展示下：
import torch
N=5
B=3
S=2
group_idx0 = torch.arange(N, dtype=torch.long)
group_idx1=group_idx0.view(1, 1, N)
group_idx2=group_idx1.repeat([B, S, 1])
print("g0:",group_idx0)
print("g1:",group_idx1)
print("g2:",group_idx2)
#结果：
g0: tensor([0, 1, 2, 3, 4])
g1: tensor([[[0, 1, 2, 3, 4]]])
g2: tensor([[[0, 1, 2, 3, 4],
[0, 1, 2, 3, 4]],
[[0, 1, 2, 3, 4],
[0, 1, 2, 3, 4]],
[[0, 1, 2, 3, 4],
[0, 1, 2, 3, 4]]])
2、对group_idx.sort的理解：
torch.sort(input, dim=-1, descending=False, out=None)，dim=-1说的是最后一维，在源码中指的就是dim=2
a=torch.randn(2,3,4)
print("a",a)
print("dim=0",torch.sort(a,0))
print("dim=1",torch.sort(a,1))
print("dim=2",torch.sort(a,2))
print("dim=-1",torch.sort(a,-1))
#结果
a tensor([[[ 0.1644, -0.9524, -0.0522, -1.7683],
[-0.0426, -1.3940, -0.9358, -2.5367],
[ 0.6171, 0.2587, 1.6798, 0.3828]],
[[ 1.0571, -0.2126, -0.1489, 0.5902],
[ 0.1673, -0.5937, -0.3240, 1.1439],
[-0.4273, -0.4449, -0.8735, -0.6969]]])
dim=0 (tensor([[[ 0.1644, -0.9524, -0.1489, -1.7683],
[-0.0426, -1.3940, -0.9358, -2.5367],
[-0.4273, -0.4449, -0.8735, -0.6969]],
[[ 1.0571, -0.2126, -0.0522, 0.5902],
[ 0.1673, -0.5937, -0.3240, 1.1439],
[ 0.6171, 0.2587, 1.6798, 0.3828]]]))
dim=1 (tensor([[[-0.0426, -1.3940, -0.9358, -2.5367],
[ 0.1644, -0.9524, -0.0522, -1.7683],
[ 0.6171, 0.2587, 1.6798, 0.3828]],
[[-0.4273, -0.5937, -0.8735, -0.6969],
[ 0.1673, -0.4449, -0.3240, 0.5902],
[ 1.0571, -0.2126, -0.1489, 1.1439]]])
dim=2 (tensor([[[-1.7683, -0.9524, -0.0522, 0.1644],
[-2.5367, -1.3940, -0.9358, -0.0426],
[ 0.2587, 0.3828, 0.6171, 1.6798]],
[[-0.2126, -0.1489, 0.5902, 1.0571],
[-0.5937, -0.3240, 0.1673, 1.1439],
[-0.8735, -0.6969, -0.4449, -0.4273]]])
dim=-1 (tensor([[[-1.7683, -0.9524, -0.0522, 0.1644],
[-2.5367, -1.3940, -0.9358, -0.0426],
[ 0.2587, 0.3828, 0.6171, 1.6798]],
[[-0.2126, -0.1489, 0.5902, 1.0571],
[-0.5937, -0.3240, 0.1673, 1.1439],
[-0.8735, -0.6969, -0.4449, -0.4273]]])
经过group_idx.sort(dim=-1)[0][:, :, :nsample]之后group_idx的维度为[B,S,nsample].
3、对group_idx[mask] = group_first[mask]的理解：
import torch
N=5
B=3
S=2
group_idx0 = torch.arange(N, dtype=torch.long)
group_idx1=group_idx0.view(1, 1, N)
group_idx2=group_idx1.repeat([B, S, 1])
mask= group_idx2 == 3
print(mask)
print(group_idx2[mask])
group_idx2[mask] =10
print(group_idx2)
#结果：
maks: tensor([[[0, 0, 0, 1, 0],
[0, 0, 0, 1, 0]],
[[0, 0, 0, 1, 0],
[0, 0, 0, 1, 0]],
[[0, 0, 0, 1, 0],
[0, 0, 0, 1, 0]]], dtype=torch.uint8)
group_idx2[mask]: tensor([3, 3, 3, 3, 3, 3])
group_idx2: tensor([[[ 0, 1, 2, 10, 4],
[ 0, 1, 2, 10, 4]],
[[ 0, 1, 2, 10, 4],
[ 0, 1, 2, 10, 4]],
[[ 0, 1, 2, 10, 4],
[ 0, 1, 2, 10, 4]]])
我们可以得出这样的结论： mask必须是一个 ByteTensor ，而且shape必须和 a一样 并且元素只能是0或者1 ，是将 mask中为1的元素所在的索引，在a中相同的的索引处替换为 value ,mask value必须同为tensor