为解决这一问题，我们从理论出发，提出了 EPro-PnP 模块，其输出位姿的概率密度分布而非单一的位姿最优解，从而将不可导的最优位姿替换为了可导的概率密度，实现了稳定的端到端训练。EPro-PnP 通用性强，适用于各类具体任务和数据，可以用于改进现有的基于 PnP 的位姿估计方法，也可以借助其灵活性训练全新的网络。从更一般的意义来说，EPro-PnP 本质是将常见的分类 softmax 带入到了连续域，理论上可以推广至训练一般的嵌套了优化层的模型。
论文链接：https://arxiv.org/abs/2203.13254
代码链接：https://github.com/tjiiv-cprg/EPro-PnP
一、前言
我们研究的是 3D 视觉中的一个经典问题：基于单张 RGB 图像定位其中的 3D 物体。具体而言，给定一张含有 3D 物体投影的图像，我们的目标是确定物体坐标系到相机坐标系的刚体变换。这一刚体变换被称为物体的位姿，记作 y，其包含两部分：1）位置（position）分量，可用 3x1 的位移向量 t 表示，2）朝向（orientation）分量，可用 3x3 的旋转矩阵 R 表示。
针对这一问题，现有方法可以分为显式和隐式两大类。显式方法也可称作直接位姿预测，即使用前馈神经网络（FFN）直接输出物体位姿的各个分量，通常是：1）预测物体的深度，2）找出物体中心点在图像上的 2D 投影位置，3）预测物体的朝向（朝向的具体处理方法可能比较复杂）。利用标有物体真实位姿的图像数据，可以设计损失函数直接监督位姿预测结果，轻松地实现网络的端到端训练。然而，这样的网络缺乏可解释性，在规模较小的数据集上易于过拟合。在 3D 目标检测任务中，显式方法占据主流，尤其是对于规模较大的数据集（例如 nuScenes）。
隐式方法则是基于几何优化的位姿估计方法，最典型的代表是基于 PnP 的位姿估计方法。这类方法中，首先需要在图像坐标系中找出 N 个 2D 点（第 i 点 2D 坐标记作），同时在物体坐标系中找出与之相关联的 N 个 3D 点（第 i 点 3D 坐标记作），有时还需要获取各对点的关联权重（第 i 对点的关联权重记作）。根据透视投影约束，这 N 对 2D-3D 加权关联点隐式地定义了物体的最优位姿。具体而言，我们可以找出使重投影误差最小的物体位姿：
其中，表示加权重投影误差，是位姿的函数。表示含有内参的相机投影函数，表示元素乘积。PnP 方法常见于物体几何形状已知的 6 自由度位姿估计任务中。
基于 PnP 的方法也需要前馈网络去预测 2D-3D 关联点集。相比于直接位姿预测，这一深度学习结合传统几何视觉算法的模型有非常好的可解释性，其泛化性能较为稳定，但在以往的工作中模型的训练方法存在缺陷。很多方法通过构建代理损失函数，去监督 X 这一中间结果，这对于位姿而言不是最优的目标。例如，已知物体形状的前提下，可以预先选取出物体的 3D 关键点，然后训练网络去找出对应的 2D 投影点位置。这也意味着代理损失只能学习 X 中的部分变量，因此不够灵活。如果我们不知道训练集中物体的形状，需要从零开始学习 X 中的全部内容该怎么办？
显示和隐式方法的优势互补，如果能够通过监督 PnP 输出的位姿结果，端到端地训练网络去学习关联点集 X ，则可以将二者优势结合。为实现这一目标，一些近期研究利用隐函数求导实现了 PnP 层的反向传播。然而，PnP 中的 argmin 函数在某些点是不连续不可导的，使得反向传播并不稳定，直接训练难以收敛。
二、EPro-PnP 方法介绍
1、EPro-PnP 模块
为了实现稳定的端到端训练，我们提出了端到端概率 PnP（end-to-end probabilistic PnP），即 EPro-PnP。其基本思想是将隐式位姿视作一个概率分布，则其概率密度对于 X 是可导的。首先基于重投影误差定义位姿的似然函数：
若使用无信息先验，则位姿的后验概率密度为似然函数的归一化结果：
可以注意到，以上公式与常用的分类 softmax 公式分接近，其实 EPro-PnP 的本质就是将softmax从离散阈搬到了连续阈，把求和换成了积分。
2、KL 散度损失
在训练模型的过程中，已知物体真实位姿，则可以定义目标位姿分布。此时可以计算 KL 散度作为训练网络所用的损失函数（因固定，也可以理解为交叉熵损失函数）。在目标趋近于 Dirac 函数的情况下，基于 KL 散度的损失函数可以简化为以下形式：
如对其求导则有：
可见，该损失函数由两项构成，第一项（记作）试图降低位姿真值的重投影误差，第二项（记作）试图增大预测位姿各处的重投影误差。二者方向相反，效果如下图（左）所示。作为类比，右边就是我们在训练分类网络时常用的分类交叉熵损失。
3、蒙特卡洛位姿损失
需要注意到，KL 损失中的第二项中含有积分，这一积分没有解析解，因此必须通过数值方法进行近似。综合考虑通用性，精确度和计算效率，我们采用蒙特卡洛方法，通过采样来模拟位姿分布。
具体而言，我们采用了一种重要性采样算法——Adaptive Multiple Importance Sampling（AMIS），计算出K个带有权重的位姿样本，我们将这一过程称作蒙特卡洛 PnP：
据此，第二项可以近似为关于权重的函数，且可以反向传播：
位姿采样的可视化效果如下图所示：
4、针对 PnP 求解器的导数正则化
尽管蒙特卡洛 PnP 损失可以用于训练网络得到高质量的位姿分布，但在推理阶段，还是需要通过 PnP 优化求解器来得到最优位姿解。常用的高斯 - 牛顿及其衍生算法通过迭代优化求解，其迭代增量是由代价函数的一阶和二阶导数决定的。为使 PnP 的解更接近真值，可以对代价函数的导数进行正则化。设计正则化损失函数如下：
其中，为高斯 - 牛顿迭代增量，与代价函数的一阶和二阶导数有关，且可以反向传播，表示距离度量，对于位置使用 smooth L1，对于朝向使用 cosine similarity。在不一致时，该损失函数促使迭代增量指向实际真值。
三、基于 EPro-PnP 的位姿估计网络
我们在 6 自由度位姿估计和 3D 目标检测两个子任务上分别使用了不同的网络。其中，对于 6 自由度位姿估计，在 ICCV 2019 的 CDPN 网络的基础上稍加修改并用 EPro-PnP 训练，用来进行 ablation studies；对于 3D 目标检测，在 ICCVW 2021 的 FCOS3D 基础上设计了全新的变形关联（deformable correspondence）检测头，以证明 EPro-PnP 可以训练网络在没有物体形状知识的情况下直接学出所有 2D-3D 点和关联权重，从而展现 EPro-PnP 在应用方面的灵活性。
1、用于 6 自由度位姿估计的稠密关联网络
网络结构如上图所示，只是在原版 CDPN 的基础上修改了输出层。原版 CDPN 使用已经检测到的物体 2D 框裁剪出区域图像，输入到 ResNet34 backbone 中。原版 CDPN 将位置与朝向解耦为两个分支，位置分支使用直接预测的显式方法，而朝向分支使用稠密关联和 PnP 的隐式方法。为了研究 EPro-PnP，改动后的网络只保留了稠密关联分支，其输出为 3 通道的 3D 坐标图，以及 2 通道关联权重，其中关联权重经过了 spatial softmax 和 global weight scaling。增加 spatial softmax 目的是对权重进行归一化，使其具有类似 attention map 的性质，可以关注相对重要的区域，实验证明权重归一化也是稳定收敛的关键。Global weight scaling 反映了位姿分布的集中程度。该网络仅需 EPro-PnP 的蒙特卡洛位姿损失就可以训练，此外可以增加导数正则化，以及在物体形状已知的情况下增加额外的 3D 坐标回归损失。
2、用于 3D 目标检测的变形关联网络
网络结构如上图所示。总体而言是基于 FCOS3D 检测器，参考 deformable DETR 设计的网络结构。在 FCOS3D 的基础上，保留其 centerness 和 classification 层，而将其原有的位姿预测层替换为 object embedding 和 reference point 层，用于生成 object query。参考 deformable DETR，我们通过预测相对于 reference point 的偏移量得到 2D 采样位置（也就得到了）。采样后的 feature 经由 attention 操作聚合为 object feature，用于预测物体级别的结果（3D score，weight scale，3D box size 等）。此外，采样后各点的 feature 在加入 object embedding 并经由 self attention 处理后输出各点所对应的的 3D 坐标和关联权重。所预测的全部可由 EPro-PnP 的蒙特卡洛位姿损失训练得到，不需要额外正则化就可以收敛并有较高的精度。在此基础上，可以增加导数正则化损失和辅助损失进一步提升精度。
四、实验结果
1、6 自由度位姿估计任务
使用 LineMOD 数据集实验，并严格与 CDPN baseline 进行比对，主要结果如上。可见，增加 EPro-PnP 损失进行端到端训练，精度显著提升（+12.70）。继续增加导数正则化损失，精度进一步提升。在此基础上，使用原版 CDPN 的训练结果初始化并增加 epoch（保持总 epoch 数与原版 CDPN 的完整三阶段训练一致）可以使精度进一步提升，其中预训练 CDPN 的优势部分来源于 CDPN 训练时有额外的 mask 监督。
上图是 EPro-PnP 与各种领先方法的比较。由较落后的 CDPN 改进而来的 EPro-PnP 在精度上接近 SOTA，并且 EPro-PnP 的架构简洁，完全基于 PnP 进行位姿估计，不需要额外进行显式深度估计或位姿精修，因此在效率上也有优势。
2、3D 目标检测任务
使用 nuScenes 数据集实验，与其他方法对比结果如上图所示。EPro-PnP 不仅相对 FCOS3D 有了明显提升，还超越了当时的 SOTA、FCOS3D 的另一个改进版本 PGD。更重要的是，EPro-PnP 目前是唯一在 nuScenes 数据集上使用几何优化方法估计位姿的。因 nuScenes 数据集规模较大，端到端训练的直接位姿估计网络已具有较好性能，而我们的结果说明了端到端地训练基于几何优化的模型能做到在大数据集上取得更加优异的性能。
3、可视化分析
上图显示了用 EPro-PnP 训练的稠密关联网络的预测结果。其中，关联权重图对图像中的重要区域进行了高光，类似于 attention 机制。由损失函数分析可知，高光区域对应的是重投影不确定性较低以及对位姿变动较为敏感的区域。
3D 目标检测的结果如上图所示。其中左上视图显示了变形关联网络采样出的 2D 点位置，红色表示水平 X 分量较高的带你，绿色表示垂直 Y 分量较高的点。绿色点一般位于物体上下两端，其主要作用是通过物体高度来推算物体的距离，这一特性并非人为指定，完全是自由训练的结果。右图显示了俯视图上的检测结果，其中蓝色云图表示物体中心点位置的分布密度，反映了物体定位的不确定性。一般远处的物体定位不确定性大于近处的物体。
EPro-PnP 的另一重要优势在于，能够通过预测复杂的多峰分布来表示朝向的模糊性。如上图所示，Barrier 由于物体本身旋转对称，朝向经常出现相差 180° 的两个峰值；Cone 本身没有特定的朝向，因此预测结果在各个方向均有分布；Pedestrian 虽不完全旋转对称，但因图像不清晰，不易判断正面和背面，有时也会出现两个峰值。这一概率特性使得 EPro-PnP 对于对称物体不需要在损失函数上做任何特殊处理。
五、总结
EPro-PnP 将原本不可导的最优位姿转变为可导的位姿概率密度，使得基于 PnP 几何优化的位姿估计网络可实现稳定且灵活的端到端训练。EPro-PnP 可应用于一般的 3D 物体位姿估计问题，即使在未知 3D 物体几何形状的情况下，也可以通过端到端训练学习得到物体的 2D-3D 关联点。因此，EPro-PnP 拓宽了网络设计的可能性，例如我们提出的变形关联网络，这在以往是不可能训练的。
此外，EPro-PnP 也可以直接被用于改进现有的基于 PnP 的位姿估计方法，通过端到端训练释放现有网络的潜力，提升位姿估计精度。从更一般的意义来说，EPro-PnP 本质是将常见的分类 softmax 带入到了连续域，不仅可用于其他基于几何优化的 3D 视觉问题，理论上还可以推广至训练一般的嵌套了优化层的模型。
