文章来源
知乎：https://zhuanlan.zhihu.com/p/289620126
Github：https://github.com/zhulf0804/PCReg.PyTorch
源码分析
源码来自Github
目录结构
PCReg.PyTorch
data
ModelNet40.py # ModelNet模型的Dataset定义
loss
earth_mover_distance # 推土机距离损失定义
models
fgr.py # Fast Global Registration (FGR) 算法执行点云之间的配准
utils
dist.py # 计算的是两组点之间点对点的平方距离
format.py # 读取pcd、npy和pcd文件相互转换
time.py # 一个装饰器，计算传入函数的执行时间
custom_evaluate.py # 评估ICP和基于迭代的网络后结果的性能
custom_train.py # 进行配准任务的训练
modelnet40_evaluate.py # 同上，只是模型换成了modelnet40
数据处理
Copy from utils import readpcd
from utils import pc_normalize, random_select_points, shift_point_cloud, \
jitter_point_cloud, generate_random_rotation_matrix, \
generate_random_tranlation_vector, transform
class CustomData(Dataset):
def __init__(self, root, npts, train=True):
super(CustomData, self).__init__()
# 根据train参数选择训练集或验证集的目录名
dirname = 'train_data' if train else 'val_data'
path = os.path.join(root, dirname)
self.train = train
# 获取所有文件路径
self.files = [os.path.join(path, item) for item in sorted(os.listdir(path))]
self.npts = npts
def __getitem__(self, item):
# 获取指定索引的数据项
file = self.files[item]
ref_cloud = readpcd(file, rtype='npy') # 加载pcd文件
ref_cloud = random_select_points(ref_cloud, m=self.npts) # 随机选点
ref_cloud = pc_normalize(ref_cloud) # 归一化
R, t = generate_random_rotation_matrix(-20, 20), \
generate_random_tranlation_vector(-0.5, 0.5)
src_cloud = transform(ref_cloud, R, t) # 随机生成旋转平移并应用
if self.train: # 如果是训练集，对点云添加随机噪声
ref_cloud = jitter_point_cloud(ref_cloud)
src_cloud = jitter_point_cloud(src_cloud)
return ref_cloud, src_cloud, R, t
def __len__(self):
return len(self.files) 初始化：CustomData 类继承自 PyTorch 的 Dataset 类，接受数据集根目录、采样点数和指示是否为训练集的参数。根据参数选择训练集或验证集的目录名，读取该目录下的所有点云文件路径。
获取数据项：__getitem__ 方法用于按索引获取指定数据项。首先从文件列表中获取对应索引的点云文件，然后加载该点云数据。
随机选点 ：对点云数据执行随机选点操作，将点云数据减少到指定的点数。这有助于减小计算复杂度并提高处理速度。
归一化 ：将点云数据进行归一化处理，将其坐标范围限制在一个固定的范围内。这有助于提高模型的泛化性能和收敛速度。
生成随机变换矩阵 ：生成一个随机的旋转矩阵和平移向量，用于模拟真实场景中点云之间的相对位置变化。旋转角度范围为 -20 到 20 度，平移范围为 -0.5 到 0.5。
应用变换：将生成的随机旋转矩阵和平移向量应用到参考点云数据，得到源点云数据。
添加随机噪声 ：如果是训练集，为参考点云和源点云添加随机噪声。这有助于提高模型的鲁棒性，使其能够在噪声环境下更好地执行点云配准。
返回数据：返回预处理后的参考点云、源点云、旋转矩阵和平移向量 。这些数据将作为网络的输入和监督信号进行训练。
数据集大小：__len__ 方法返回数据集中点云文件的数量。
网络定义
Encoder (PointNet)
Copy class PointNet(nn.Module): # 用于提取点云特征
def __init__(self, in_dim, gn, mlps=[64, 64, 64, 128, 1024]):
super(PointNet, self).__init__()
self.backbone = nn.Sequential()
# 初始化网络的骨架（backbone），它是一个 Sequential 容器，用于按顺序存放神经网络的层
# 遍历 mlps 列表中的每个输出维度（out_dim）
for i, out_dim in enumerate(mlps):
# 向骨架中添加 1D 卷积层，输入维度为 in_dim，输出维度为 out_dim
self.backbone.add_module(f'pointnet_conv_{i}',
nn.Conv1d(in_dim, out_dim, 1, 1, 0))
# 如果 gn（GroupNorm）为 True，则添加 GroupNorm 层
if gn:
self.backbone.add_module(f'pointnet_gn_{i}',
nn.GroupNorm(8, out_dim))
# 向骨架中添加 ReLU 激活函数
self.backbone.add_module(f'pointnet_relu_{i}',
nn.ReLU(inplace=True))
in_dim = out_dim
def forward(self, x):
x = self.backbone(x)
x, _ = torch.max(x, dim=2)
return x 这个网络用于提取点云数据的特征，为一个1024维的特征
输入维度：in_dim (例如 3，表示位置信息)
全局最大池化（沿 num_points 维度）：维度变为 (batch_size, 1024)
在每个卷积层之后，我们都添加了一个 ReLU 激活函数，以引入非线性，使得网络能够学习更复杂的特征表示。这种网络结构可以更有效地捕捉点云数据中的信息。在全局最大池化之后，我们得到一个尺寸为 (batch_size, 1024) 的特征向量，可以用于后续任务。
Decoder (Benchmark)
Copy class Benchmark(nn.Module): # 这个名字其实不是很恰当，PointCloudAlignment
def __init__(self, gn, in_dim1, in_dim2=2048, fcs=[1024, 1024, 512, 512, 256, 7]):
super(Benchmark, self).__init__()
self.in_dim1 = in_dim1 # 输入维度
self.encoder = PointNet(in_dim=in_dim1, gn=gn) # 编码：PointNet提取特征
self.decoder = nn.Sequential() # 解码网络
for i, out_dim in enumerate(fcs):
self.decoder.add_module(f'fc_{i}', nn.Linear(in_dim2, out_dim)) # 添加线性层
if out_dim != 7:
if gn: # 是否添加GroupNorm层
self.decoder.add_module(f'gn_{i}',nn.GroupNorm(8, out_dim))
self.decoder.add_module(f'relu_{i}', nn.ReLU(inplace=True))
in_dim2 = out_dim
def forward(self, x, y):
# 使用编码器对输入 x 和 y 进行特征提取
x_f, y_f = self.encoder(x), self.encoder(y)
# 沿着维度 1 对特征 x_f 和 y_f 进行拼接
concat = torch.cat((x_f, y_f), dim=1)
# 将拼接后的特征传入解码器
out = self.decoder(concat)
# 对输出结果进行分割，得到平移向量（batch_t）和四元数（batch_quat）
batch_t, batch_quat = out[:, :3], out[:, 3:] / torch.norm(out[:, 3:], dim=1, keepdim=True)
# 将四元数转换为旋转矩阵（batch_R）
batch_R = batch_quat2mat(batch_quat)
# 根据输入的维度（in_dim1）选择不同的变换方式
if self.in_dim1 == 3:
transformed_x = batch_transform(x.permute(0, 2, 1).contiguous(),
batch_R, batch_t)
elif self.in_dim1 == 6:
transformed_pts = batch_transform(x.permute(0, 2, 1)[:, :, :3].contiguous(),
batch_R, batch_t)
transformed_nls = batch_transform(x.permute(0, 2, 1)[:, :, 3:].contiguous(),
batch_R)
transformed_x = torch.cat([transformed_pts, transformed_nls], dim=-1)
else:
raise ValueError
return batch_R, batch_t, transformed_x 输入点云 x 和 y 的维度：(batch_size, in_dim1, num_points)。
使用 PointNet 编码器对输入 x 和 y 进行特征提取：
输出特征维度：(batch_size, 1024)。
将特征 x_f 和 y_f 进行拼接：
拼接后的特征维度：(batch_size, 2048)。
将拼接后的特征传入解码器。解码器包含一系列全连接层和 ReLU 激活函数：
输出：(batch_size, 7)，其中前 3 个元素表示平移向量，后 4 个元素表示旋转四元数。
Benchmark 网络首先使用 PointNet 编码器对输入点云 x 和 y 提取特征，然后将两个特征向量拼接在一起。接下来，通过一个解码器（由一系列全连接层和 ReLU 激活函数组成）将拼接后的特征映射到旋转和平移参数。最后，网络输出一个大小为 (batch_size, 7) 的向量，表示点云之间的相对变换。
IterativeNet (IterativeBenchmark)
Copy class IterativeBenchmark(nn.Module): # 迭代
def __init__(self, in_dim, niters, gn):
super(IterativeBenchmark, self).__init__()
# 初始化 Benchmark 模型
self.benckmark = Benchmark(gn=gn, in_dim1=in_dim)
# 设置迭代次数
self.niters = niters
def forward(self, x, y):
# 用于存储每次迭代后变换的 x
transformed_xs = []
device = x.device
B = x.size()[0]
# 创建一个与 x 相同的张量 transformed_x
transformed_x = torch.clone(x)
# 初始化旋转矩阵为单位矩阵，大小为 Bx3x3
batch_R_res = torch.eye(3).to(device).unsqueeze(0).repeat(B, 1, 1)
# 初始化平移向量为零矩阵，大小为 Bx3x1
batch_t_res = torch.zeros(3, 1).to(device).unsqueeze(0).repeat(B, 1, 1)
# 进行 self.niters 次迭代
for i in range(self.niters):
# 对 transformed_x 和 y 使用 Benchmark 模型
batch_R, batch_t, transformed_x = self.benckmark(transformed_x, y)
# 将本次迭代的变换后的 x 添加到列表中
transformed_xs.append(transformed_x)
# 更新旋转矩阵和平移向量
batch_R_res = torch.matmul(batch_R, batch_R_res)
batch_t_res = torch.matmul(batch_R, batch_t_res) \
+ torch.unsqueeze(batch_t, -1)
# 调整 transformed_x 的维度顺序
transformed_x = transformed_x.permute(0, 2, 1).contiguous()
batch_t_res = torch.squeeze(batch_t_res, dim=-1)
#transformed_x = transformed_x.permute(0, 2, 1).contiguous()
return batch_R_res, batch_t_res, transformed_xsn 一个基于迭代方法的点云对其网络，用Pointnet作为编码器，Benchmark作为解码器，多次迭代来优化点云之间的变换
以下是 IterativeBenchmark 网络的各层和维度变化：
输入点云 x 和 y 的维度：(batch_size, in_dim, num_points)。
初始化变换后的点云 transformed_x 为与输入点云 x 相同。
初始化旋转矩阵为单位矩阵，大小为 (batch_size, 3, 3)。
初始化平移向量为零矩阵，大小为 (batch_size, 3, 1)。
对于 i 从 0 到 niters-1 进行迭代：
使用 Benchmark 模型对 transformed_x 和 y 进行特征提取，计算旋转矩阵 batch_R 和平移向量 batch_t，并获得变换后的点云 transformed_x。
调整变换后的点云 transformed_x 的维度顺序。
输出：最终的旋转矩阵 batch_R_res，最终的平移向量 batch_t_res 和每次迭代后变换的点云 transformed_xs。
IterativeBenchmark 网络首先初始化一个 Benchmark 网络，然后对输入点云 x 和 y 进行多次迭代。在每次迭代中，网络计算 transformed_x 和 y 之间的相对变换，并更新旋转矩阵和平移向量。经过 niters 次迭代后，网络输出最终的旋转矩阵、平移向量以及每次迭代后的变换点云。这种迭代策略有助于提高网络的性能和鲁棒性。