[CVPR 2020] RPM-Net: Robust Point Matching using Learned Features

零、概要

• 论文: RPM-Net: Robust Point Matching using Learned Features
• tag: CVPR 2020; registration
• 代码: https://github.com/yewzijian/RPMNet/
• 作者: Zi Jian Yew, Gim Hee Lee
• 机构: Department of Computer Science, National University of Singapore
• 笔者整理了一个最近几年150多篇点云的论文列表，欢迎大家一块学习交流。

0.0 摘要

ICP算法通过迭代的方式解决点云的刚性配准: (1) 基于最近空间距离寻找hard的点对应关系; (2) 求解最小二乘变换。其中第(1)步对待配准点云的初始位置和噪声点/离群点敏感，会导致迭代收敛到错误的局部最小值。在本文中，作者提出了RPM-Net，一种对初始化不太敏感且鲁棒的、基于学习的用于刚性点云配准方法。为此，RMP-Net使用可微的Sinkhorn层和退火(Annealing)技术，利用从空间坐标和局部几何信息中学习到的混合特征，实现点对应(point correspondences)的软(soft)分配。为了提高配准性能，作者引入了第二个网络来预测最佳退火参数。与现有的一些方法不同，RPM-Net可以处理部分可见和丢失correspondences的点云。实验结果表明，与现有的非深度学习和最近的基于深度学习的方法相比，RPM-Net达到了SOTA。

读到这里，我们会产生如下的疑问:

• RPM-Net怎么解决部分可见和丢失correspondences的点云配准 ?
• Sinkhorn和退火是什么，它们如何产生点对应的软分配的 ?

上述问题将会在二、论文的方法中介绍。

一、论文的出发点和贡献

现有点云配准的方案主要包括以下三种:

• 基于距离的匹配算法: ICP, RPM等。这类算法对待配准点云的初始位置或噪声敏感。
• 基于特征的匹配算法: 基于特征(LRF、FPH、FPFH等)建立point correspondences。这类算法要求点云具有独特的几何结构。
• 基于学习的匹配算法: PointNetLK, DCP, PRNet等。上述算法或者不能解决局部-局部点云的匹配，或者设计太复杂。

针对上述点云配准中的现状，作者提出了RPM-Net，它是一个end-to-end的可微分深度网络，它保留了RPM对噪声/离群点的鲁棒性，同时从学习的特征距离(而不是空间距离)中预测点对应(point correspondences)关系，降低了配准对初始值的依赖。

它和ICP算法类似，也是一个迭代配准算法。由于使用了混合特征(空间特征和集合特征)，算法可以在少量迭代中收敛。

在我看来，作者的主要贡献在于: (1)引进了match matrix，以及如何生成match matrix; (2)提出了一种改进的Chamfer distance度量，以提高点云局部可见情况下的配准质量。

二、论文的方法

RPM-Net的架构如Figure 2所示, 它是一个迭代网络。在每次迭代中，Source点云X和Reference点云Y会经过特征提取模块，参数预测模块，计算Match Matrix模块和SVD分解模块。当前迭代求解的R, t，会作用于当前的source点云X，并将得到的点云与Reference点云Y作为输入传递到下一次迭代中。

2.1 特征提取模块

对于点云中的一个点${\text{x}}_c$，定义它的邻域$N({\text{x}}_c)$，则

$$F_{{\text{x}}_c}=f_{\theta }({\text{x}}_c,\{\Delta {\text{x}}_{c,i}\},\{PPF({\text{x}}_c,{\text{x}}_i)\})$$

$f_{\theta }$是一个神经网络，待学习的参数是$\theta$，${\text{x}}_i\in N({\text{x}}_c)$，$\Delta {\text{x}}_{c,i}={\text{x}}_i-{\text{x}}_c$，$PPF({\text{x}}_c,{\text{x}}_i)$是一个4D的特征，其组成如下:

$$PPF({\text{x}}_c,{\text{x}}_i)=(\angle ({\text{n}}_c,\Delta {\text{x}}_{c,i}),\angle ({\text{n}}_i,\Delta {\text{x}}_{c,i}),\angle ({\text{n}}_c,{\text{n}}_i),||\Delta {\text{x}}_{c,i}|{|}_2)$$

$\angle$表示角度angle(v1, v2) = atan2(cross(v1, v2), dot(v1, v2))。PPF特征是一个四维向量，因此对于特定的$i$, $({\text{x}}_c,\{\Delta {\text{x}}_{c,i}\},\{PPF({\text{x}}_c,{\text{x}}_i)\})$是一个10维向量。

特征提取模块如Figure 2(b)所示，对于${\text{x}}_c$的每个邻域点${\text{x}}_i$，构成一个10维向量，batch的数据格式为(B, 10, N), 送入到一个给定共享参数的MLP，得到维度为(B, C1, N)的特征, 而后经过Max Pooling，得到维度为(B, C1)的特征，接着经过几个MLP，得到维度为(B, C2)的特征。最后对特征进行归一化，产生维度为(B, C2)的特征$F_{{\text{x}}_c}$。

2.2 参数预测模块

在原始的RPM算法中，参数$\alpha ,\beta$依赖于数据集，需要人工调制。在RMP-Net中，这些参数由于依赖于学习到的特征，很难调制。因此，作者设计了一个网络来预测这两个参数。如Figure 2©所示，网络的输入来自source点云X(J, 3)和reference点云Y(K, 3)，组织形式为concat, 构成维度为(J + K, 3)的输入。另外为了表征当前点来自于哪个数据集，增加了第4列特征，0表示当前点来自于X点云, 1表示当前点来自于Y点云。所以，参数预测模块的输入数据是(J+K, 4)维的。对于(B, 4, J + K)的数据首先经过共享参数的MLP，产生维度为(B, C3, J+K)的数据，经过Max Pooling，得到维度为(B, C3)的数据，接下来经过MLP得到(B, 2)的数据。为了确保预测的$\alpha ,\beta$是正的，作者在最后使用了softplux激活函数。

2.3 计算Match Matrix

经过上述两个模块，得到了两个特征矩阵$F_X\in {\mathbb{R}}^{J\times C},F_Y\in {\mathbb{R}}^{K\times C}$和退火参数$\alpha ,\beta$。接下来如何找点对应(point correspondences)关系呢 ? 这里是通过Match Matrix $M\in {\mathbb{R}}^{J\times K}$实现的。$M$的初始值$m_{jk}$通过如下方式获得:

$$m_{jk}=e^{-\beta (||F_{{\text{x}}_j}-F_{{\text{y}}_k}|{|}_2-\alpha )}$$

接下来是论文的一个重点，如Figure 2(d)所示, 通过Sinkhorn迭代来获取一个双随机限制的矩阵(doubly stochastic matrix)。相关的理论证明在A relationship between arbitrary positive matrices and doubly stochastic matrices。这里把Sinkhorn相关代码搬过来吧。它是一个迭代算法，大致做法交替进行行归一化和列归一化。

def sinkhorn(log_alpha, n_iters: int = 5, slack: bool = True, eps: float = -1) -> torch.Tensor:
    """ Run sinkhorn iterations to generate a near doubly stochastic matrix, where each row or column sum to <=1

    Args:
        log_alpha: log of positive matrix to apply sinkhorn normalization (B, J, K)
        n_iters (int): Number of normalization iterations
        slack (bool): Whether to include slack row and column
        eps: eps for early termination (Used only for handcrafted RPM). Set to negative to disable.

    Returns:
        log(perm_matrix): Doubly stochastic matrix (B, J, K)

    Modified from original source taken from:
        Learning Latent Permutations with Gumbel-Sinkhorn Networks
        https://github.com/HeddaCohenIndelman/Learning-Gumbel-Sinkhorn-Permutations-w-Pytorch
    """

    # Sinkhorn iterations
    prev_alpha = None
    if slack:
        zero_pad = nn.ZeroPad2d((0, 1, 0, 1))
        log_alpha_padded = zero_pad(log_alpha[:, None, :, :])

        log_alpha_padded = torch.squeeze(log_alpha_padded, dim=1)

        for i in range(n_iters):
            # Row normalization
            log_alpha_padded = torch.cat((
                    log_alpha_padded[:, :-1, :] - (torch.logsumexp(log_alpha_padded[:, :-1, :], dim=2, keepdim=True)),
                    log_alpha_padded[:, -1, None, :]),  # Don't normalize last row
                dim=1)

            # Column normalization
            log_alpha_padded = torch.cat((
                    log_alpha_padded[:, :, :-1] - (torch.logsumexp(log_alpha_padded[:, :, :-1], dim=1, keepdim=True)),
                    log_alpha_padded[:, :, -1, None]),  # Don't normalize last column
                dim=2)

            if eps > 0:
                if prev_alpha is not None:
                    abs_dev = torch.abs(torch.exp(log_alpha_padded[:, :-1, :-1]) - prev_alpha)
                    if torch.max(torch.sum(abs_dev, dim=[1, 2])) < eps:
                        break
                prev_alpha = torch.exp(log_alpha_padded[:, :-1, :-1]).clone()

        log_alpha = log_alpha_padded[:, :-1, :-1]
    else:
        for i in range(n_iters):
            # Row normalization (i.e. each row sum to 1)
            log_alpha = log_alpha - (torch.logsumexp(log_alpha, dim=2, keepdim=True))

            # Column normalization (i.e. each column sum to 1)
            log_alpha = log_alpha - (torch.logsumexp(log_alpha, dim=1, keepdim=True))

            if eps > 0:
                if prev_alpha is not None:
                    abs_dev = torch.abs(torch.exp(log_alpha) - prev_alpha)
                    if torch.max(torch.sum(abs_dev, dim=[1, 2])) < eps:
                        break
                prev_alpha = torch.exp(log_alpha).clone()

    return log_alpha

这里就回答了Sinkhorn和退火是什么, 如何产生point soft correspondence将在2.4小节介绍。

先解释一下RPM-Net怎么解决部分可见和丢失correspondences的点云配准的吧。

因为X中的每一个点${\text{x}}_j$并不总是在Y中存在一个对应点${\text{y}}_k$，因此$M$矩阵的第$j$行之和就是0，不满足doubly stochastic 矩阵的性质。因此，作者在$M$矩阵最右边增加了一列，这样就可以满足doubly stochastic的性质了；同理在$M$矩阵的最下面增加了一行。这样就产生了$M\in {\mathbb{R}}^{(J+1)\times (K+1)}$的矩阵了。在行归一化时，不归一最后一行，在列归一化时不归一化最后一列。结束后，取前$J$行, 前$M$列，带入下面的计算。

2.4 计算R, t

对于点云X中每一个点${\text{x}}_j$，构造响应的对应点:

$${\text{y}}_j=\frac{1}{{\Sigma }_{k=1}^K{m}_{j,k}}{\Sigma }_{k=1}^K{m}_{j,k}{\text{y}}_k$$

另外，因为并不是每一个${\text{x}}_j$都有一个对应点，因此在每一个对应关系$({\text{x}}_j,{\text{y}}_i)$设置一个权重$w_j={\Sigma }_{k=1}^K{m}_{j,k}$。最后通过SVD求解R, t。

对每一个对应关系$({\text{x}}_j,{\text{y}}_i)$设置一个权重$w_j$是否由必要呢 ?

应该是必要的，论文中对每一个${\text{x}}_j$构造了soft correspondence, 虽然soft correspondence是根据权重生成的，但它仍然是不准确的(比如, 权重为0, 生成了(0, 0, 0)，此时的对应关系就是错误的)，因此在最后的对应关系有必要增加一个权重，来惩罚不同置信度下correspondence所带来的损失。

2.5 损失函数

$$L_{\text{reg}}=\frac{1}{J}{\Sigma }_j^J|(R_{\text{gt}}{\text{x}}_j+t_{\text{gt}})-(R_{\text{pred}}{\text{x}}_j+t_{\text{pred}})|$$

这个损失函数感觉还是挺有创意的，既考虑了R, t，还考虑点云X中的所有点。

作者说，只使用上述loss，网络倾向于把大部分点标记为outliers。为此，作者增加了第2个loss，鼓励$M$得到更大的inliers:

$$L_{\text{inlier}}=-\frac{1}{J}{\Sigma }_j^J{\Sigma }_k^K{m}_{jk}-\frac{1}{K}{\Sigma }_k^K{\Sigma }_j^J{m}_{jk}$$

整体的loss为两者的和:

$$L_{\text{total}}=L_{\text{reg}}+\lambda L_{\text{inlier}}$$

因为RPM-Net是一个迭代的网络，因此考虑每次迭代的损失，作者这里把损失设置为每次迭代损失的加权和:

$$L={\Sigma }_i^N(\frac{1}{2}{)}^{(N-i)}{L}_{\text{total}}^i$$

还有一个需要注意的地方: 每次迭代都是独立的，都是学习GroundTruth的R,t。因此，作者在训练时采用了2次迭代，在测试时采用了5次迭代。这个和以往的一些算法，如PRNet是不同的。

三、论文的实验

作者在ModelNet40上进行了实验，使用前20类进行训练和验证(调参)，使用其余20类进行测试。训练集有5112个，验证集有1202个，测试集有1266个。在训练和测试时，旋转角度在[0, 45°]进行采样，平移从[-0.5, 0.5]进行采样。

3.1评价指标

• anisotropic

  旋转矩阵和平移向量的MSE和MAE值，在DCP和PRNet中都有使用。

• isotropic
  $$\text{Error}(R)=\text{arc cos}(\frac{\text{tr}(R_{\text{GT}}^{-1}R)-1}{2}),\text{Error}(t)=||t_{\text{GT}}-t|{|}_2$$

  此方式在3DRegNet论文中有使用。

• Chamfer distance

  anisotropic和isotropic的评价都是在R,t上面， Chamfer distance有些不同，它考虑了样本。在计算最近距离时，考虑了原始点云，并不是只比较ref和src之间的距离。

  $$CD(X,Y)=\frac{1}{|X|}{\Sigma }_{x\in X}\underset{y\in Y_{\text{clean}}}{\min }||x-y|{|}_2+\frac{1}{|Y|}{\Sigma }_{y\in Y}\underset{x\in X_{\text{clean}}}{\min }||x-y|{|}_2$$

作者在Clean, Gaussian Noise, Partial Visibility数据集上进行了实验，并与ICP, RPM, FGR, PointNetLK, DCP-v2算法进行了比较，可视化结果如Figure 3所示。我对这部分的实验细节比较感兴趣，所以记录一下作者是怎么处理这些实验的(大部分是结合作者开源的代码进行描述)。

3.2 Clean Data

• 训练集: 对于单个点云S(2048, ), 随机采样1024个点, 生成src和ref点云, 这里src和ref是相同的点云。随机生成R, t, 将R, t作用于src得到新的点云src_transformed。随机打乱src_transformed和ref点云中点的顺序，并分别记这两个点云为src, ref。此时训练的数据和Gt就是src, ref和R, t。
• 测试集: 对于单个点云S(2048, ), 固定采样1024个点, 生成src和ref点云, 这里src和ref是相同的点云。接下来的过程和训练集一致。

实验结果如Table 2所示，RPM-Net和FGR处于第一或者第二的位置。在接下来的章节中，可以看懂FGRd对噪声非常敏感。

3.3 Gaussian Noise

• 训练集: 对于单个点云S(2048, )，首先生成src和ref点云, 这两个点云是一致的且都包含2048个点。随机生成R, t，将R, t作用于src得到新的点云src_transformed。从src_transformed和ref点云随机采样1024个点生成新的点云, 得到src_transformed_sampled和ref_sampled, 此时得到的src_transformed_sampled和ref_sampled不具有一一对应关系的。接下来对src_transformed_sampled和ref_sampled进行随机抖动和随机打乱顺序，生成新的点云，分别即为src, ref。
• 测试集: 生成过程和训练集一样，只是设定了随机种子。

实验结果如Table 2所示，RPM-Net优于其他方法。

3.4 Partial Visibility

• 训练集: 对于单个点云S(2048, )，首先生成src和ref点云, 这两个点云是一致的且都包含2048个点。对src和ref进行partial visibility(这里通过随机一个方向的半空间，移动确保剩余的点在原始点云的70%左右)点云生成新的点云src_crop和ref_crop，这两个点云是不一样的。接下来就是随机生成R,t等操作，和Gaussian Noise的数据处理方式相同。
• 测试集: 成过程和训练集一样，只是设定了随机种子。

实验结果如Table 3所示, RPM-Net明显优于其它方法。

3.5 Ablation Studies

作者在partial visibility实验上进行了ablation studies，来更好的了解各种选择是如何影响算法的性能的，实验结果如Table 4所示。比较1-3行和第5行，可以看到x, $\Delta \text{x}$, PPF都是必须的，且x对网络的性能影响较大。第4行，把使用复杂网络预测$\alpha ,\beta$改成了学习$\alpha ,\beta$，相比于第5行，性能也降低了。

四、对论文的想法

• 缺点
  • 只在ModelNet40上进行了实验，不具有说服力。
  • 只比较了基于学习方法的PointNetLK和DCP(局部点云匹配不佳)，比较不充分。
• 优点
  • RPM-Net是一个迭代网络，在设计损失函数和累加每次迭代损失时具有创新性；而且，每次迭代都是独立的，都是学习GroundTruth的R,t，这一点与以往的算法是不同的。
  • 引入传统的RPM和SkinHorn的思想，设计基于深度学习的配准方法。

五、RPM-Net预测可视化代码

在src文件夹中新建vis.py，内容如下:

"""Evaluate RPMNet. Also contains functionality to compute evaluation metrics given transforms

Example Usages:
    1. Visualize RPMNet
        python vis.py --noise_type crop --resume [path-to-model.pth] --dataset_path [your_path]/modelnet40_ply_hdf5_2048

"""
import os
import open3d as o3d
import random
from tqdm import tqdm
import torch

from arguments import rpmnet_eval_arguments
from common.misc import prepare_logger
from common.torch import dict_all_to_device, CheckPointManager, to_numpy
from common.math_torch import se3
from data_loader.datasets import get_test_datasets
import models.rpmnet


def vis(npys):
    pcds = []
    colors = [[1.0, 0, 0],
              [0, 1.0, 0],
              [0, 0, 1.0]]
    for ind, npy in enumerate(npys):
        color = colors[ind] if ind < 3 else [random.random() for _ in range(3)]
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(npy)
        pcd.paint_uniform_color(color)
        pcds.append(pcd)
    return pcds


def inference_vis(data_loader, model: torch.nn.Module):
    _logger.info('Starting inference...')
    model.eval()

    with torch.no_grad():
        for data in tqdm(data_loader):

            #gt_transforms = data['transform_gt']
            points_src = data['points_src'][..., :3]
            points_ref = data['points_ref'][..., :3]
            #points_raw = data['points_raw'][..., :3]
            dict_all_to_device(data, _device)
            pred_transforms, endpoints = model(data, _args.num_reg_iter)
            src_transformed = se3.transform(pred_transforms[-1], points_src)

            src_np = torch.squeeze(points_src).cpu().detach()
            src_transformed_np = torch.squeeze(src_transformed).cpu().detach()
            ref_np = torch.squeeze(points_ref).cpu().detach()

            pcds = vis([src_np, src_transformed_np, ref_np])
            o3d.visualization.draw_geometries(pcds)


def get_model():
    _logger.info('Computing transforms using {}'.format(_args.method))
    assert _args.resume is not None
    model = models.rpmnet.get_model(_args)
    model.to(_device)
    if _device == torch.device('cpu'):
        model.load_state_dict(
            torch.load(_args.resume, map_location=torch.device('cpu'))['state_dict'])
    else:
        model.load_state_dict(torch.load(_args.resume)['state_dict'])
    return model


def main():
    # Load data_loader
    test_dataset = get_test_datasets(_args)
    test_loader = torch.utils.data.DataLoader(test_dataset,
                                              batch_size=1, shuffle=False)

    model = get_model()
    inference_vis(test_loader, model)  # Feedforward transforms

    _logger.info('Finished')


if __name__ == '__main__':
    # Arguments and logging
    parser = rpmnet_eval_arguments()
    _args = parser.parse_args()
    _logger, _log_path = prepare_logger(_args, log_path=_args.eval_save_path)
    os.environ['CUDA_VISIBLE_DEVICES'] = str(_args.gpu)
    if _args.gpu >= 0 and (_args.method == 'rpm' or _args.method == 'rpmnet'):
        os.environ['CUDA_VISIBLE_DEVICES'] = str(_args.gpu)
        _device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
    else:
        _device = torch.device('cpu')

    main()

部分-部分匹配的可视化结果如下:

红色为src点云，蓝色为ref点云，绿色为RPM-Net输出R,t作用到src的点云。

使用方法:

python vis.py --noise_type crop --resume [path-to-model.pth] --dataset_path [your_path]/modelnet40_ply_hdf5_2048

如,
# python vis.py --noise_type crop --resume partial-trained.pth --dataset_path /Users/zhulf/data/modelnet40_ply_hdf5_2048