3D点云配准RPM-Net模型解读（附论文+源码）

论文链接：RPM-Net: Robust Point Matching using Learned Features

官方链接：RPMNet

老规矩，先看看效果。

看看论文里给的对比图

总体流程

在学之前得储备PointNet++的知识，因为RPM-Net是通过PointNet++提取输入特征的。因为之前讲过Pointnet++，所以这里不会重复讲PointNet++的内容。

RPMNet是一个用于点云配准的模型，基于传统算法RPM做的，而RPM是非常吃经验值$\alpha$和$\beta$的，RPMNet通过神经网络来预测这两个值。

算法的总体框架如下图。

（b）这个特征提取模块和PointNet++是一样，不过这里有些不同的是，这里会提取10个特征。首先xyz、dxyz肯定是有的，然后还有4个就是这个PPF。PointNet++ 里不是有多簇嘛，每个簇里有很多个点，这些点与簇的关系是怎么样的呢？这就是PPF要干的事情，最后会得到4个关系特征。

（c）X，Y分别是原点云数据和参考点云数据（原点云数据仅是在原数据上采样了，或者裁剪了，位置什么的都没有变化。参考点云数据在采样裁剪的基础上还会有一些扰动，比如旋转平移之类的）。然后经过一系列的Conv1D卷积，maxpool，连几个MLP全连接层，最后激活输出$\alpha$和$\beta$。

（d）这里输入的$F_x$和$F_y$是(b)的输出，$\alpha$和$\beta$是©的输出。最后我们得到啥？变换的矩阵啊。这个变换矩阵的的求解论文到没详细说，只是引用了两篇参考文献，表示照着这个人家这个做的，有兴趣的可以去看看这两篇文献[DeepICP,Deep Closest Point]。

代码

数据预处理

数据预处理这块很简单，在路径src/data_loader/datasets.py的方法get_train_datasets。

def get_train_datasets(args: argparse.Namespace):train_categories, val_categories = None, Noneif args.train_categoryfile:  # 训练类别train_categories = [line.rstrip('\n') for line in open(args.train_categoryfile)]train_categories.sort()if args.val_categoryfile:  # 验证类别val_categories = [line.rstrip('\n') for line in open(args.val_categoryfile)]val_categories.sort()# 根据指定的参数对数据进行变换train_transforms, val_transforms = get_transforms(args.noise_type, args.rot_mag, args.trans_mag,args.num_points, args.partial)_logger.info('Train transforms: {}'.format(', '.join([type(t).__name__ for t in train_transforms])))_logger.info('Val transforms: {}'.format(', '.join([type(t).__name__ for t in val_transforms])))train_transforms = torchvision.transforms.Compose(train_transforms)val_transforms = torchvision.transforms.Compose(val_transforms)if args.dataset_type == 'modelnet_hdf':  # 从数据集中选择了部分数据train_data = ModelNetHdf(args.dataset_path, subset='train', categories=train_categories,transform=train_transforms)val_data = ModelNetHdf(args.dataset_path, subset='test', categories=val_categories,transform=val_transforms)else:raise NotImplementedErrorreturn train_data, val_data

步骤很简单吧，首先拿到数据的类别信息，然后根据参数对数据进行变换。数据集里并不是所有的数据都适合做这个任务的，所以作者选了一部分的数据来做。主要看看get_transforms吧，别的也没啥好说的。

def get_transforms(noise_type: str,rot_mag: float = 45.0, trans_mag: float = 0.5,num_points: int = 1024, partial_p_keep: List = None):partial_p_keep = partial_p_keep if partial_p_keep is not None else [0.7, 0.7]  # 裁剪比例if noise_type == "clean":...elif noise_type == "jitter":...elif noise_type == "crop":# Both source and reference point clouds cropped, plus same noise in "jitter"train_transforms = [Transforms.SplitSourceRef(),  # 复制一份 源点云和参考点云Transforms.RandomCrop(partial_p_keep),  # 裁剪Transforms.RandomTransformSE3_euler(rot_mag=rot_mag, trans_mag=trans_mag),  # 旋转 平移Transforms.Resampler(num_points),  # 采样num_points个点Transforms.RandomJitter(),  # 高斯噪声Transforms.ShufflePoints()]  # 打乱点的顺序test_transforms = [Transforms.SetDeterministic(),Transforms.SplitSourceRef(),Transforms.RandomCrop(partial_p_keep),Transforms.RandomTransformSE3_euler(rot_mag=rot_mag, trans_mag=trans_mag),Transforms.Resampler(num_points),Transforms.RandomJitter(),Transforms.ShufflePoints()]else:raise NotImplementedErrorreturn train_transforms, test_transforms

因为clean和jitter的内容都包含在crop，我就直接说crop了。作者没有直接用transforms.Compose的方法哈，所有的变换都是自己写的，在路径src/data_loader/transforms.py下。我就不去transforms.py里细讲了，这个看名字也知道干啥了。SplitSourceRef()拷贝一份点云数据（源点云和参考点云）、RandomCrop()随机裁剪、RandomTransformSE3_euler()随机旋转和平移、Resampler()随机采样、RandomJitter()添加高斯噪声、ShufflePoints()随机打乱。

模型

主要来看看模型的前向传播是怎么搞的，在路径src/models/rpmnet.py里RPMNet的方法forward。

def forward(self, data, num_iter: int = 1):endpoints = {}xyz_ref, norm_ref = data['points_ref'][:, :, :3], data['points_ref'][:, :, 3:6]xyz_src, norm_src = data['points_src'][:, :, :3], data['points_src'][:, :, 3:6]xyz_src_t, norm_src_t = xyz_src, norm_src  # (b,717,3) 717=1024*0.7transforms = []all_gamma, all_perm_matrices, all_weighted_ref = [], [], []all_beta, all_alpha = [], []for i in range(num_iter):beta, alpha = self.weights_net([xyz_src_t, xyz_ref])feat_src = self.feat_extractor(xyz_src_t, norm_src_t)  # (b,717,96)feat_ref = self.feat_extractor(xyz_ref, norm_ref)  # (b,717,96)feat_distance = match_features(feat_src, feat_ref)  # 计算距离 (b,717,717)affinity = self.compute_affinity(beta, feat_distance, alpha=alpha)  # (b,717,717)# Compute weighted coordinateslog_perm_matrix = sinkhorn(affinity, n_iters=self.num_sk_iter, slack=self.add_slack)perm_matrix = torch.exp(log_perm_matrix)weighted_ref = perm_matrix @ xyz_ref / (torch.sum(perm_matrix, dim=2, keepdim=True) + _EPS)  # (b,717,3)# Compute transform and transform pointstransform = compute_rigid_transform(xyz_src, weighted_ref, weights=torch.sum(perm_matrix, dim=2))  # (b,3,4)xyz_src_t, norm_src_t = se3.transform(transform.detach(), xyz_src, norm_src)transforms.append(transform)all_gamma.append(torch.exp(affinity))all_perm_matrices.append(perm_matrix)all_weighted_ref.append(weighted_ref)all_beta.append(to_numpy(beta))all_alpha.append(to_numpy(alpha))endpoints['perm_matrices_init'] = all_gammaendpoints['perm_matrices'] = all_perm_matricesendpoints['weighted_ref'] = all_weighted_refendpoints['beta'] = np.stack(all_beta, axis=0)endpoints['alpha'] = np.stack(all_alpha, axis=0)return transforms, endpoints

先将源点和参考点的xyz坐标与法向量坐标拿出来吗，方便后面做计算。好，直接看for循环。之前有说过传统算法RPM是非常吃经验值的，$\alpha$和$\beta$需要人为估计一个值。作者一想，好麻烦，直接使用神经网络来预测一个吧。来看看self.weights_net怎么做的。

计算$\alpha$和$\beta$

在路径src/models/feature_nets.py里ParameterPredictionNet的方法forward。

def forward(self, x):src_padded = F.pad(x[0], (0, 1), mode='constant', value=0)  # 多了一个标识符 src的为0ref_padded = F.pad(x[1], (0, 1), mode='constant', value=1)  # ref的为1  (b,717,4) 4->xyz+标识符concatenated = torch.cat([src_padded, ref_padded], dim=1)  # (b,1434,4)prepool_feat = self.prepool(concatenated.permute(0, 2, 1))  # (b,4,1434) -> (b,1024,1434)pooled = torch.flatten(self.pooling(prepool_feat), start_dim=-2)  # (b,1024)raw_weights = self.postpool(pooled)  # (b,2)beta = F.softplus(raw_weights[:, 0])alpha = F.softplus(raw_weights[:, 1])return beta, alpha

首先将源点云数据和参考点云数据拼接在一起，在此之前呢会给它们添个标识符，源的为0，参考的为1。

这里解释一下为什么经常使用permute来转变维度，因为pytorch里channel first，特征维度要放在第一位（也就是batch后的第一位）。正常特征图的维度是(c,h,w)，但这里只有两个维度啊，2D卷积用不了，所以用的是一维卷积Conv1d。看一下self.prepool的结构。

self.prepool = nn.Sequential(nn.Conv1d(4, 64, 1),nn.GroupNorm(8, 64),nn.ReLU(),nn.Conv1d(64, 64, 1),nn.GroupNorm(8, 64),nn.ReLU(),nn.Conv1d(64, 64, 1),nn.GroupNorm(8, 64),nn.ReLU(),nn.Conv1d(64, 128, 1),nn.GroupNorm(8, 128),nn.ReLU(),nn.Conv1d(128, 1024, 1),nn.GroupNorm(16, 1024),nn.ReLU(),
)

self.prepool里有好几层Conv1d，可以看到channel的的变化，4->64->64->64->128->1024。这里的nn.GroupNorm也是一种归一化的方式，第一个参数是num_groups，表示将channel划分为num_groups组，比如第一层输出是(8, 64, 1434)，会将64个channel划分为8组，每组8个channel。然后每个样本独立计算（不依赖 batch 统计信息），每个组共享均值和方差，然后归一化。（像常见的BN，会依赖batch size，因为它是计算整个 batch 维度的均值和方差）。

然后做一个maxpooling，将数据浓缩成精华。比如pooling之前，1434个点有1024个特征；pooling完之后，不是每个点有啥特征了，而是全局的一个特征是一个1024维向量。现在做完这些我们得有一个输出了，还记得我们初衷吗？预测$\alpha$和$\beta$啊。看一下self.postpool的结构。

self.postpool = nn.Sequential(nn.Linear(1024, 512),nn.GroupNorm(16, 512),nn.ReLU(),nn.Linear(512, 256),nn.GroupNorm(16, 256),nn.ReLU(),nn.Linear(256, 2 + np.prod(weights_dim)),
)

连了几个全连接，可以看到维度的转变， 1024->512->256->2。输出的结果可以看到有正有负，不过$\alpha$和$\beta$的值需要是正数，所以连个F.softplus，可以看做激活函数，最后输出。

=======================================================================

回到RPMNet的forward，现在需要对原点云数据和参考点云数据进行特征提取，来看看self.feat_extractor是怎么做的。

特征提取

在路径src/models/feature_nets.py里FeatExtractionEarlyFusion的方法forward。先展示部分。

def forward(self, xyz, normals):features = sample_and_group_multi(-1, self.radius, self.n_sample, xyz, normals)features['xyz'] = features['xyz'][:, :, None, :]

之前PointNet++里也有sample_and_group_multi，代码一模一样的，不过这里面添了个PPF，来看看怎么回事。

def sample_and_group_multi(npoint: int, radius: float, nsample: int, xyz: torch.Tensor, normals: torch.Tensor,returnfps: bool = False):B, N, C = xyz.shapeif npoint > 0:S = npointfps_idx = farthest_point_sample(xyz, npoint)  # [B, npoint, C]new_xyz = index_points(xyz, fps_idx)nr = index_points(normals, fps_idx)[:, :, None, :]else:S = xyz.shape[1]fps_idx = torch.arange(0, xyz.shape[1])[None, ...].repeat(xyz.shape[0], 1).to(xyz.device)new_xyz = xyznr = normals[:, :, None, :]  # (b,717,1,3)idx = query_ball_point(radius, nsample, xyz, new_xyz, fps_idx)  # (B, npoint, nsample)grouped_xyz = index_points(xyz, idx)  # (B, npoint, nsample, C)  npoint个簇 每个簇nsample个点 每个点C个特征d = grouped_xyz - new_xyz.view(B, S, 1, C)  # d = p_r - p_i  (B, npoint, nsample, 3) 去均值(邻域点相对中心点的位置偏移量)ni = index_points(normals, idx)nr_d = angle(nr, d)  # 中心点法向量与邻域点坐标偏移量的夹角  (b,717,64)ni_d = angle(ni, d)  # 邻域点法向量与邻域点坐标偏移量的夹角  (b,717,64)nr_ni = angle(nr, ni)  # 中心点法向量与邻域点法向量的夹角  (b,717,64)d_norm = torch.norm(d, dim=-1)   # 邻域点相对中心点的距离  (b,717,64)xyz_feat = d  # (B, npoint, n_sample, 3)ppf_feat = torch.stack([nr_d, ni_d, nr_ni, d_norm], dim=-1)  # (B, npoint, n_sample, 4)if returnfps:return {'xyz': new_xyz, 'dxyz': xyz_feat, 'ppf': ppf_feat}, grouped_xyz, fps_idxelse:return {'xyz': new_xyz, 'dxyz': xyz_feat, 'ppf': ppf_feat}

这里输入的npoint为-1，表示每个点都作为中心点画圈采样。query_ball_point之前在PointNet++里面讲的很详细了，这里就不说了。得到的grouped_xyz的size为(B, 717, 64, 3)，可以理解为：以所有点为中心画圈，圈成717个簇，每个簇64个点，每个点3个特征。

nr_d的size为(B,717,64)，这个计算的是中心点法向量与邻域点坐标偏移量的夹角，可以理解为：每簇里有64个点，这64个点与对应簇的关系。同理，ni_d、nr_ni和d_norm你都可以这样理解，就是计算簇里的点与簇的关系，论文里也没解释为啥这样做。

=======================================================================

回到FeatExtractionEarlyFusion的forward中，继续往下看。

# Gate and concat
concat = []
for i in range(len(self.features)):  # 为了将10个特征拼在一起f = self.features[i]expanded = (features[f]).expand(-1, -1, self.n_sample, -1)concat.append(expanded)
fused_input_feat = torch.cat(concat, -1)  # (B,717,64,10)# Prepool_FC, pool, postpool-FC
new_feat = fused_input_feat.permute(0, 3, 2, 1)  # [B, 10, n_sample, N]
new_feat = self.prepool(new_feat)pooled_feat = torch.max(new_feat, 2)[0]  # Max pooling (B, C, N)post_feat = self.postpool(pooled_feat)  # Post pooling dense layers
cluster_feat = post_feat.permute(0, 2, 1)
cluster_feat = cluster_feat / torch.norm(cluster_feat, dim=-1, keepdim=True)return cluster_feat  # (B, N, C)

for循环是为了将xyz、dxyz和ppf的特征拼接在一起，得到fused_input_feat的size为(b,717,64,10)。ok。现在我们有3个维度了，可以做2D卷积了，来看看self.prepool的结构。

Sequential(nn.Conv2d(10, 96, 1),nn.GroupNorm(8, 96),nn.ReLU(),nn.Conv2d(96, 96, 1),nn.GroupNorm(8, 96),nn.ReLU(),nn.Conv2d(96, 192, 1),nn.GroupNorm(8, 192),nn.ReLU(),
)

可以看到维度的变化，10->96->96->192。然后接着Max pooling得到size为(b,192,717)。这下又只有两个维度了，只能用1D卷积了，来看一下self.postpool的结构。

Sequential(nn.Conv1d(192, 192, 1),nn.GroupNorm(8, 192),nn.ReLU(),nn.Conv1d(192, 96, 1),nn.GroupNorm(8, 96),nn.ReLU(),nn.Conv1d(96, 96, 1),
)

可以看到维度的变化，192->96。最后进行L2归一化，得到最终输出，size为(B,717,96)

。

=======================================================================

回到RPMNet的forward，继续往下看。现在要计算变换矩阵了。

变换矩阵

1、计算加权坐标

现在源点云数据和参考点云数据的特征都提取出来了，match_features来计算它俩之间的距离。每个源点与所有的参考点计算距离（相似度），构成一个J × K 的匹配矩阵（这里J、K是一样的），size为(b,717,717)。compute_affinity的解释是计算匹配矩阵值的对数， 看一下是怎么计算的。

def compute_affinity(self, beta, feat_distance, alpha=0.5):if isinstance(alpha, float):hybrid_affinity = -beta[:, None, None] * (feat_distance - alpha)else:hybrid_affinity = -beta[:, None, None] * (feat_distance - alpha[:, None, None])return hybrid_affinity

很简单啊，公式就是$\beta \times (匹配矩阵-\alpha )$，这个公式可以理解为，衡量每对点之间是否是对应点（源点-参考点 对应关系）。你要问我为什么这样就可以衡量？那我包支支吾吾的，应该涉及传统RPM算法的数学公式吧…。不过，很显然这里面也没计算对数啊，我看了一下在下面的sinkhorn计算了。

sinkhorn的解释是将输入矩阵归一化为近似的双重随机矩阵（每行、每列的和小于等于1）。就不跳进去看了，代码很简单。结果出来又进行了torch.exp，刚log完又exp回来，就理解为将刚刚的矩阵归一化了就行，得到源点和参考点的匹配关系。最后计算加权的参考点坐标，用刚刚的概率矩阵作为权重（因为归一化了嘛，就当概率矩阵看吧），与参考点进行加权求和，最终得到每个源点对应的加权参考点坐标，然后归一化。

（注：这里应该是涉及了一些数学公式的推导，但是论文里并没有详细说，就这么理解吧）

2、计算刚体变换矩阵和变换点

compute_rigid_transform是计算刚体变换的，即找到一个 旋转矩阵和平移向量 ，使得源点变换后最接近参考点。看一下compute_rigid_transform干了啥吧。

def compute_rigid_transform(a: torch.Tensor, b: torch.Tensor, weights: torch.Tensor):weights_normalized = weights[..., None] / (torch.sum(weights[..., None], dim=1, keepdim=True) + _EPS)  # (b,717,1)centroid_a = torch.sum(a * weights_normalized, dim=1)  # (b,3)centroid_b = torch.sum(b * weights_normalized, dim=1)  # (b,3)a_centered = a - centroid_a[:, None, :]  # 去中心化 (b,717,3)b_centered = b - centroid_b[:, None, :]  # (b,717,3)cov = a_centered.transpose(-2, -1) @ (b_centered * weights_normalized)  # (b,3,3)# Compute rotation using Kabsch algorithm. Will compute two copies with +/-V[:,:3]# and choose based on determinant to avoid flipsu, s, v = torch.svd(cov, some=False, compute_uv=True)rot_mat_pos = v @ u.transpose(-1, -2)v_neg = v.clone()v_neg[:, :, 2] *= -1rot_mat_neg = v_neg @ u.transpose(-1, -2)rot_mat = torch.where(torch.det(rot_mat_pos)[:, None, None] > 0, rot_mat_pos, rot_mat_neg)assert torch.all(torch.det(rot_mat) > 0)# Compute translation (uncenter centroid)translation = -rot_mat @ centroid_a[:, :, None] + centroid_b[:, :, None]transform = torch.cat((rot_mat, translation), dim=2)return transform

先将权重归一化，这个权重weights就是我刚刚说的概率矩阵的和。对源点坐标和参考点坐标进行去中心化，使其绕 (0,0,0) 旋转。然后计算它俩之间的协方差矩阵，…，woc下面全是数学公式，看得我头疼。

下面代码看不懂，看注释上写的是计算旋转矩阵、处理镜像翻转、计算平移向量，最终找到源点到参考点的最优刚体变换（旋转+平移）。返回的是刚体变换矩阵size为(b,3,4)，3表示xyz，4：前三列是旋转矩阵，最后一列是平移向量。

======================================================================

回到RPMNet的forward，继续往下看。

这时候我们已经有变换矩阵了，那得看看这个矩阵效果怎么样啊。se3.transform将源点通过变换矩阵得到变换后的点，用于下一次迭代的特征提取和匹配。将每次循环的结果都存储起来，为后续损失计算做准备。

计算损失

来到路径src/train.py下看看compute_losses损失是什么计算的。

def compute_losses(data: Dict, pred_transforms: List, endpoints: Dict,loss_type: str = 'mae', reduction: str = 'mean') -> Dict:losses = {}num_iter = len(pred_transforms)# Compute lossesgt_src_transformed = se3.transform(data['transform_gt'], data['points_src'][..., :3])  # (b,717,3)if loss_type == 'mse':# MSE loss to the groundtruth (does not take into account possible symmetries)criterion = nn.MSELoss(reduction=reduction)for i in range(num_iter):pred_src_transformed = se3.transform(pred_transforms[i], data['points_src'][..., :3])if reduction.lower() == 'mean':losses['mse_{}'.format(i)] = criterion(pred_src_transformed, gt_src_transformed)elif reduction.lower() == 'none':losses['mse_{}'.format(i)] = torch.mean(criterion(pred_src_transformed, gt_src_transformed),dim=[-1, -2])elif loss_type == 'mae':# MSE loss to the groundtruth (does not take into account possible symmetries)criterion = nn.L1Loss(reduction=reduction)for i in range(num_iter):pred_src_transformed = se3.transform(pred_transforms[i], data['points_src'][..., :3])if reduction.lower() == 'mean':losses['mae_{}'.format(i)] = criterion(pred_src_transformed, gt_src_transformed)elif reduction.lower() == 'none':losses['mae_{}'.format(i)] = torch.mean(criterion(pred_src_transformed, gt_src_transformed),dim=[-1, -2])else:raise NotImplementedError# Penalize outliersfor i in range(num_iter):ref_outliers_strength = (1.0 - torch.sum(endpoints['perm_matrices'][i], dim=1)) * _args.wt_inlierssrc_outliers_strength = (1.0 - torch.sum(endpoints['perm_matrices'][i], dim=2)) * _args.wt_inliersif reduction.lower() == 'mean':losses['outlier_{}'.format(i)] = torch.mean(ref_outliers_strength) + torch.mean(src_outliers_strength)elif reduction.lower() == 'none':losses['outlier_{}'.format(i)] = torch.mean(ref_outliers_strength, dim=1) + \torch.mean(src_outliers_strength, dim=1)discount_factor = 0.5  # Early iterations will be discountedtotal_losses = []for k in losses:discount = discount_factor ** (num_iter - int(k[k.rfind('_')+1:]) - 1)total_losses.append(losses[k] * discount)losses['total'] = torch.sum(torch.stack(total_losses), dim=0)return losses

又看到se3.transform老朋友了，先将源点通过变换矩阵变换过来，这里用的都是真实值，是真货。判断一下用哪种损失函数，然后又看到se3.transform兄弟，这次还是源点通过变换矩阵变换过来，but，这次用的是我们预测出来的变换矩阵，是假货。ok，让我们来看看真货与假货之间的差异。

单单计算配准损失，不够不够，得再来点离群惩罚。上面说过，perm_matrices里存放的是源点与参考点的匹配概率，对行或列求和再取1-sum，就是匹配失败的程度。

之前在rpmnet里计算时，不是迭代的计算了几次，每次迭代的数据都被存储下来了。很显然，前面迭代的时候只是粗糙的计算了一下，后面再此基础上再微调，让其越来越准确。因此，需要加个权重discount ，前面迭代的结果损失轻点，但是到后面了，就不能一直这么碌碌无为了，得加重点。最终加权叠加所有损失返回结果。