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

為解決這一問題，我們從理論出發，提出了 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 視覺問題，理論上還可以推廣至訓練一般的嵌套了優化層的模型。