本文解讀我們獲得 CVPR 2022 最佳學生論文獎的工作《EPro-PnP: Generalized End-to-End Probabilistic Perspective-n-Points for Monocular Object Pose Estimation》。論文研究的問題是基于單張圖像估計物體在 3D 空間中的位姿?，F(xiàn)有方法中，基于 PnP 幾何優(yōu)化的位姿估計方法往往通過深度網(wǎng)絡提取 2D-3D 關聯(lián)點，然而因為位姿最優(yōu)解在反向傳播時存在不可導的問題，難以實現(xiàn)以位姿誤差作為損失對網(wǎng)絡進行穩(wěn)定的端到端訓練，此時 2D-3D 關聯(lián)點依賴其他代理損失的監(jiān)督，這對于位姿估計而言不是最佳的訓練目標。

為解決這一問題，我們從理論出發(fā)，提出了 EPro-PnP 模塊，其輸出位姿的概率密度分布而非單一的位姿最優(yōu)解，從而將不可導的最優(yōu)位姿替換為了可導的概率密度，實現(xiàn)了穩(wěn)定的端到端訓練。EPro-PnP 通用性強，適用于各類具體任務和數(shù)據(jù)，可以用于改進現(xiàn)有的基于 PnP 的位姿估計方法，也可以借助其靈活性訓練全新的網(wǎng)絡。從更一般的意義來說，EPro-PnP 本質是將常見的分類 softmax 帶入到了連續(xù)域，理論上可以推廣至訓練一般的嵌套了優(yōu)化層的模型。

論文鏈接：https://arxiv.org/abs/2203.13254

代碼鏈接：https://github.com/tjiiv-cprg/EPro-PnP

一、前言

我們研究的是 3D 視覺中的一個經(jīng)典問題：基于單張 RGB 圖像定位其中的 3D 物體。具體而言，給定一張含有 3D 物體投影的圖像，我們的目標是確定物體坐標系到相機坐標系的剛體變換。這一剛體變換被稱為物體的位姿，記作 y，其包含兩部分：1）位置（position）分量，可用 3x1 的位移向量 t 表示，2）朝向（orientation）分量，可用 3x3 的旋轉矩陣 R 表示。 

針對這一問題，現(xiàn)有方法可以分為顯式和隱式兩大類。顯式方法也可稱作直接位姿預測，即使用前饋神經(jīng)網(wǎng)絡（FFN）直接輸出物體位姿的各個分量，通常是：1）預測物體的深度，2）找出物體中心點在圖像上的 2D 投影位置，3）預測物體的朝向（朝向的具體處理方法可能比較復雜）。利用標有物體真實位姿的圖像數(shù)據(jù)，可以設計損失函數(shù)直接監(jiān)督位姿預測結果，輕松地實現(xiàn)網(wǎng)絡的端到端訓練。然而，這樣的網(wǎng)絡缺乏可解釋性，在規(guī)模較小的數(shù)據(jù)集上易于過擬合。在 3D 目標檢測任務中，顯式方法占據(jù)主流，尤其是對于規(guī)模較大的數(shù)據(jù)集（例如 nuScenes）。

隱式方法則是基于幾何優(yōu)化的位姿估計方法，最典型的代表是基于 PnP 的位姿估計方法。這類方法中，首先需要在圖像坐標系中找出 N 個 2D 點（第 i 點 2D 坐標記作），同時在物體坐標系中找出與之相關聯(lián)的 N 個 3D 點（第 i 點 3D 坐標記作），有時還需要獲取各對點的關聯(lián)權重（第 i 對點的關聯(lián)權重記作）。根據(jù)透視投影約束，這 N 對 2D-3D 加權關聯(lián)點隱式地定義了物體的最優(yōu)位姿。具體而言，我們可以找出使重投影誤差最小的物體位姿：

其中，表示加權重投影誤差，是位姿的函數(shù)。表示含有內參的相機投影函數(shù)，表示元素乘積。PnP 方法常見于物體幾何形狀已知的 6 自由度位姿估計任務中。

基于 PnP 的方法也需要前饋網(wǎng)絡去預測 2D-3D 關聯(lián)點集。相比于直接位姿預測，這一深度學習結合傳統(tǒng)幾何視覺算法的模型有非常好的可解釋性，其泛化性能較為穩(wěn)定，但在以往的工作中模型的訓練方法存在缺陷。很多方法通過構建代理損失函數(shù)，去監(jiān)督 X 這一中間結果，這對于位姿而言不是最優(yōu)的目標。例如，已知物體形狀的前提下，可以預先選取出物體的 3D 關鍵點，然后訓練網(wǎng)絡去找出對應的 2D 投影點位置。這也意味著代理損失只能學習 X 中的部分變量，因此不夠靈活。如果我們不知道訓練集中物體的形狀，需要從零開始學習 X 中的全部內容該怎么辦？

顯示和隱式方法的優(yōu)勢互補，如果能夠通過監(jiān)督 PnP 輸出的位姿結果，端到端地訓練網(wǎng)絡去學習關聯(lián)點集 X ，則可以將二者優(yōu)勢結合。為實現(xiàn)這一目標，一些近期研究利用隱函數(shù)求導實現(xiàn)了 PnP 層的反向傳播。然而，PnP 中的 argmin 函數(shù)在某些點是不連續(xù)不可導的，使得反向傳播并不穩(wěn)定，直接訓練難以收斂。

二、EPro-PnP 方法介紹

1、EPro-PnP 模塊

為了實現(xiàn)穩(wěn)定的端到端訓練，我們提出了端到端概率 PnP（end-to-end probabilistic PnP），即 EPro-PnP。其基本思想是將隱式位姿視作一個概率分布，則其概率密度對于 X 是可導的。首先基于重投影誤差定義位姿的似然函數(shù)：

若使用無信息先驗，則位姿的后驗概率密度為似然函數(shù)的歸一化結果：

可以注意到，以上公式與常用的分類 softmax 公式分接近，其實 EPro-PnP 的本質就是將softmax從離散閾搬到了連續(xù)閾，把求和換成了積分。

2、KL 散度損失

在訓練模型的過程中，已知物體真實位姿，則可以定義目標位姿分布。此時可以計算 KL 散度作為訓練網(wǎng)絡所用的損失函數(shù)（因固定，也可以理解為交叉熵損失函數(shù)）。在目標趨近于 Dirac 函數(shù)的情況下，基于 KL 散度的損失函數(shù)可以簡化為以下形式：

如對其求導則有：

可見，該損失函數(shù)由兩項構成，第一項（記作）試圖降低位姿真值的重投影誤差，第二項（記作）試圖增大預測位姿各處的重投影誤差。二者方向相反，效果如下圖（左）所示。作為類比，右邊就是我們在訓練分類網(wǎng)絡時常用的分類交叉熵損失。

3、蒙特卡洛位姿損失

需要注意到，KL 損失中的第二項中含有積分，這一積分沒有解析解，因此必須通過數(shù)值方法進行近似。綜合考慮通用性，精確度和計算效率，我們采用蒙特卡洛方法，通過采樣來模擬位姿分布。

具體而言，我們采用了一種重要性采樣算法——Adaptive Multiple Importance Sampling（AMIS），計算出K個帶有權重的位姿樣本，我們將這一過程稱作蒙特卡洛 PnP：

據(jù)此，第二項可以近似為關于權重的函數(shù)，且可以反向傳播：

位姿采樣的可視化效果如下圖所示：

4、針對 PnP 求解器的導數(shù)正則化

盡管蒙特卡洛 PnP 損失可以用于訓練網(wǎng)絡得到高質量的位姿分布，但在推理階段，還是需要通過 PnP 優(yōu)化求解器來得到最優(yōu)位姿解。常用的高斯 - 牛頓及其衍生算法通過迭代優(yōu)化求解，其迭代增量是由代價函數(shù)的一階和二階導數(shù)決定的。為使 PnP 的解更接近真值，可以對代價函數(shù)的導數(shù)進行正則化。設計正則化損失函數(shù)如下：

其中，為高斯 - 牛頓迭代增量，與代價函數(shù)的一階和二階導數(shù)有關，且可以反向傳播，表示距離度量，對于位置使用 smooth L1，對于朝向使用 cosine similarity。在不一致時，該損失函數(shù)促使迭代增量指向實際真值。

三、基于 EPro-PnP 的位姿估計網(wǎng)絡

我們在 6 自由度位姿估計和 3D 目標檢測兩個子任務上分別使用了不同的網(wǎng)絡。其中，對于 6 自由度位姿估計，在 ICCV 2019 的 CDPN 網(wǎng)絡的基礎上稍加修改并用 EPro-PnP 訓練，用來進行 ablation studies；對于 3D 目標檢測，在 ICCVW 2021 的 FCOS3D 基礎上設計了全新的變形關聯(lián)（deformable correspondence）檢測頭，以證明 EPro-PnP 可以訓練網(wǎng)絡在沒有物體形狀知識的情況下直接學出所有 2D-3D 點和關聯(lián)權重，從而展現(xiàn) EPro-PnP 在應用方面的靈活性。

1、用于 6 自由度位姿估計的稠密關聯(lián)網(wǎng)絡

網(wǎng)絡結構如上圖所示，只是在原版 CDPN 的基礎上修改了輸出層。原版 CDPN 使用已經(jīng)檢測到的物體 2D 框裁剪出區(qū)域圖像，輸入到 ResNet34 backbone 中。原版 CDPN 將位置與朝向解耦為兩個分支，位置分支使用直接預測的顯式方法，而朝向分支使用稠密關聯(lián)和 PnP 的隱式方法。為了研究 EPro-PnP，改動后的網(wǎng)絡只保留了稠密關聯(lián)分支，其輸出為 3 通道的 3D 坐標圖，以及 2 通道關聯(lián)權重，其中關聯(lián)權重經(jīng)過了 spatial softmax 和 global weight scaling。增加 spatial softmax 目的是對權重進行歸一化，使其具有類似 attention map 的性質，可以關注相對重要的區(qū)域，實驗證明權重歸一化也是穩(wěn)定收斂的關鍵。Global weight scaling 反映了位姿分布的集中程度。該網(wǎng)絡僅需 EPro-PnP 的蒙特卡洛位姿損失就可以訓練，此外可以增加導數(shù)正則化，以及在物體形狀已知的情況下增加額外的 3D 坐標回歸損失。

2、用于 3D 目標檢測的變形關聯(lián)網(wǎng)絡

網(wǎng)絡結構如上圖所示?？傮w而言是基于 FCOS3D 檢測器，參考 deformable DETR 設計的網(wǎng)絡結構。在 FCOS3D 的基礎上，保留其 centerness 和 classification 層，而將其原有的位姿預測層替換為 object embedding 和 reference point 層，用于生成 object query。參考 deformable DETR，我們通過預測相對于 reference point 的偏移量得到 2D 采樣位置（也就得到了）。采樣后的 feature 經(jīng)由 attention 操作聚合為 object feature，用于預測物體級別的結果（3D score，weight scale，3D box size 等）。此外，采樣后各點的 feature 在加入 object embedding 并經(jīng)由 self attention 處理后輸出各點所對應的的 3D 坐標和關聯(lián)權重。所預測的全部可由 EPro-PnP 的蒙特卡洛位姿損失訓練得到，不需要額外正則化就可以收斂并有較高的精度。在此基礎上，可以增加導數(shù)正則化損失和輔助損失進一步提升精度。

四、實驗結果

1、6 自由度位姿估計任務

使用 LineMOD 數(shù)據(jù)集實驗，并嚴格與 CDPN baseline 進行比對，主要結果如上?？梢姡黾?EPro-PnP 損失進行端到端訓練，精度顯著提升（+12.70）。繼續(xù)增加導數(shù)正則化損失，精度進一步提升。在此基礎上，使用原版 CDPN 的訓練結果初始化并增加 epoch（保持總 epoch 數(shù)與原版 CDPN 的完整三階段訓練一致）可以使精度進一步提升，其中預訓練 CDPN 的優(yōu)勢部分來源于 CDPN 訓練時有額外的 mask 監(jiān)督。

上圖是 EPro-PnP 與各種領先方法的比較。由較落后的 CDPN 改進而來的 EPro-PnP 在精度上接近 SOTA，并且 EPro-PnP 的架構簡潔，完全基于 PnP 進行位姿估計，不需要額外進行顯式深度估計或位姿精修，因此在效率上也有優(yōu)勢。

2、3D 目標檢測任務

使用 nuScenes 數(shù)據(jù)集實驗，與其他方法對比結果如上圖所示。EPro-PnP 不僅相對 FCOS3D 有了明顯提升，還超越了當時的 SOTA、FCOS3D 的另一個改進版本 PGD。更重要的是，EPro-PnP 目前是唯一在 nuScenes 數(shù)據(jù)集上使用幾何優(yōu)化方法估計位姿的。因 nuScenes 數(shù)據(jù)集規(guī)模較大，端到端訓練的直接位姿估計網(wǎng)絡已具有較好性能，而我們的結果說明了端到端地訓練基于幾何優(yōu)化的模型能做到在大數(shù)據(jù)集上取得更加優(yōu)異的性能。

3、可視化分析

上圖顯示了用 EPro-PnP 訓練的稠密關聯(lián)網(wǎng)絡的預測結果。其中，關聯(lián)權重圖對圖像中的重要區(qū)域進行了高光，類似于 attention 機制。由損失函數(shù)分析可知，高光區(qū)域對應的是重投影不確定性較低以及對位姿變動較為敏感的區(qū)域。

3D 目標檢測的結果如上圖所示。其中左上視圖顯示了變形關聯(lián)網(wǎng)絡采樣出的 2D 點位置，紅色表示水平 X 分量較高的帶你，綠色表示垂直 Y 分量較高的點。綠色點一般位于物體上下兩端，其主要作用是通過物體高度來推算物體的距離，這一特性并非人為指定，完全是自由訓練的結果。右圖顯示了俯視圖上的檢測結果，其中藍色云圖表示物體中心點位置的分布密度，反映了物體定位的不確定性。一般遠處的物體定位不確定性大于近處的物體。

EPro-PnP 的另一重要優(yōu)勢在于，能夠通過預測復雜的多峰分布來表示朝向的模糊性。如上圖所示，Barrier 由于物體本身旋轉對稱，朝向經(jīng)常出現(xiàn)相差 180° 的兩個峰值；Cone 本身沒有特定的朝向，因此預測結果在各個方向均有分布；Pedestrian 雖不完全旋轉對稱，但因圖像不清晰，不易判斷正面和背面，有時也會出現(xiàn)兩個峰值。這一概率特性使得 EPro-PnP 對于對稱物體不需要在損失函數(shù)上做任何特殊處理。

五、總結

EPro-PnP 將原本不可導的最優(yōu)位姿轉變?yōu)榭蓪У奈蛔烁怕拭芏?，使得基?PnP 幾何優(yōu)化的位姿估計網(wǎng)絡可實現(xiàn)穩(wěn)定且靈活的端到端訓練。EPro-PnP 可應用于一般的 3D 物體位姿估計問題，即使在未知 3D 物體幾何形狀的情況下，也可以通過端到端訓練學習得到物體的 2D-3D 關聯(lián)點。因此，EPro-PnP 拓寬了網(wǎng)絡設計的可能性，例如我們提出的變形關聯(lián)網(wǎng)絡，這在以往是不可能訓練的。

此外，EPro-PnP 也可以直接被用于改進現(xiàn)有的基于 PnP 的位姿估計方法，通過端到端訓練釋放現(xiàn)有網(wǎng)絡的潛力，提升位姿估計精度。從更一般的意義來說，EPro-PnP 本質是將常見的分類 softmax 帶入到了連續(xù)域，不僅可用于其他基于幾何優(yōu)化的 3D 視覺問題，理論上還可以推廣至訓練一般的嵌套了優(yōu)化層的模型。