論文閱讀-RaftStereo

文章目錄

1 概述
2 模塊說明
- 2.1 特征抽取器
- 2.2 相關金字塔
- 2.3 多級更新算子
- 2.4 Slow-Fast GRU
- 2.5 監督
3 效果

1 概述

在雙目立體匹配中，基于迭代的模型是一種比較主流的方法，而其鼻祖就是本文要講的RaftStereo。
先來說下什么是雙目立體匹配。給定極線矯正后的左圖和右圖 $I_L, I_R)$ ，雙目立體匹配的目的就是估計一張視差圖 $d$ ， $d$ 中每個像素位置的視差值 $d_{xy}$ 表示在左圖 $I_L$ 的 $(x, y)$ 位置的像素與右圖 $I_R$ 中的 $x - d_{xy}, y)$ 位置的像素在真實3D世界中對應于同一個物理點。換句話說， $d$ 表示了左圖和右圖在水平方向上的匹配關系。
有了視差圖 $d$ ，左右相機的焦距 $f_x$ 和基線 $B$ ，我們就可以通過式(1-1)得到左圖的深度圖 $D$ ，這部分具體可以參考這篇。
$\frac{Bf}{d} \tag{1-1}$
RaftStereo的輸入是左右圖 $I_L, I_R)$ ，輸出是視差 $d$ 。由于視差 $d$ 本身與相機參數無關，因此該方式訓練得到的模型可以應用于不同參數的相機。
RaftStereo參考了Raft，由特征抽取器，相關金字塔和基于GRU的迭代更新算子組成，可見圖1-1。
RAFT整體網絡結構

圖1-1 RAFT整體網絡結構

2 模塊說明

2.1 特征抽取器

如圖1-1所示，特征抽取器指的就是encoder，共有兩種encoder，一個叫做Feature Encoder，另一個叫做Context Encoder。
Feature Encoder同時作用于左、右圖像，將每幅圖像映射為密集特征圖，進而構建相關體。該網絡由一系列殘差塊和下采樣層組成，根據實驗中使用的下采樣層數不同，可生成分辨率為輸入圖像1/4或1/8且包含256個通道的特征圖。在Feature Encoder采用了Instance Normalization。
Context Encoder的架構與特征編碼器完全相同，不同之處在于使用Batch Normalization替代了Instance Normalization，并且僅在左圖上應用Context Encoder。上下文特征用于初始化更新算子的隱藏狀態，并在每次迭代更新算子時輸入GRU中。
關于feature encoder使用IN，context encoder使用BN的疑問。有人在github的issue上面也問了這個問題，可見https://github.com/princeton-vl/RAFT-Stereo/issues/17。
feature encoder中的IN比BN更有意義，因為作者希望僅從單個圖像中導出每組相關特征。context encoder也可以用IN，但是BN更考慮整個數據集的均值和方差。
有一個細節，作者的代碼中對BN進行了freeze_bn的操作，對應的代碼是m.eval()這只是將均值和方差限制為[0, 1]，對于 $γ\gamma$ 和 $β\beta$ 并沒有進行限制， $γ\gamma$ 和 $β\beta$ 還是會進行梯度更新的。回顧一下BN的公式：
$\frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta$

2.2 相關金字塔

（1）相關體
采用特征向量之間的點積作為視覺相似度的衡量指標，對具有相同y坐標的像素進行相關計算。給定分別從 $I_L$ 和 $I_R$ 提取的特征圖 $f, g ∈ R^{H×W×D}$ ，將內積運算限制在具有相同第一索引的特征向量之間進行，可得到相關體，如下式(2-1)所示。
$Cijk=∑hfijh?gikh,C∈RH×W×W(2-1)C_{ijk} = \sum_h f_{ijh} \cdot g_{ikh}, C \in R^{H \times W \times W} \tag{2-1}$
式(2-1)看不明白的話，可以先不看 $i$ ，也就是我們只看 $y$ 方向確定的情況下，同一行的相關體計算方法。很容易看出 $C_{jk}$ 就表示左圖 $j$ 位置與右圖 $k$ 位置之間的相關性，這個相關性就是長度為 $D$ 的向量 $f_{ij}$ 和 $g_{ik}$ 的點積。
相關體的計算可以通過單個矩陣乘法來高效實現，這可以在GPU上輕松計算，并且只占用總運行時間的一小部分。
由于模型輸入是極線校正后的圖像，因此可以假設所有視差均為正值。因此，相關體的計算實際上只需針對正視差進行。然而，完整體計算的優勢在于其運算可通過矩陣乘法實現，而該運算經過高度優化。這簡化了整體架構設計，使作者能夠采用通用運算而非定制GPU內核。
（2）相關金字塔
通過重復對最后一維進行平均池化，構建了一個四層結構的相關體金字塔。金字塔的第 $k$ 層是通過使用核尺寸為2、步長為2的一維平均池化操作得到，從第 $k$ 層的體積生成的相關體 $C_{k+1}$ ，其維度為 $H×W×W/2^k$ 。雖然金字塔的每一層都具有更大的感受野，但通過僅對最后一維進行池化處理，保留了原始圖像中的高分辨率信息，從而能夠恢復極其精細的結構特征。
（3）相關查找
為了在相關金字塔中建立索引，作者定義了一個類似于Raft中定義的查找算子 $L_C$ 。給定當前視差估計值 $d$ ，作者會圍繞該估計值構建一個帶有整數偏移量的一維網格（如圖2-1所示）。這個網格用于從相關金字塔的每一層進行索引操作。
相關查找示意圖

圖2-1 相關查找示意圖

由于網格的值都是整數，在每個相關體取索引時使用了線性插值。取回的值最終被拼接到一個特征圖中。這部分的代碼可以參考下面的x0變量的計算方式。

class CorrBlock1D:def __call__(self, coords):r = self.radiuscoords = coords[:, :1].permute(0, 2, 3, 1)batch, h1, w1, _ = coords.shapeout_pyramid = []for i in range(self.num_levels):corr = self.corr_pyramid[i]dx = torch.linspace(-r, r, 2*r+1)dx = dx.view(2*r+1, 1).to(coords.device)x0 = dx + coords.reshape(batch*h1*w1, 1, 1, 1) / 2**iy0 = torch.zeros_like(x0)coords_lvl = torch.cat([x0,y0], dim=-1)corr = bilinear_sampler(corr, coords_lvl)corr = corr.view(batch, h1, w1, -1)out_pyramid.append(corr)out = torch.cat(out_pyramid, dim=-1)return out.permute(0, 3, 1, 2).contiguous().float()

2.3 多級更新算子

從初始點 $d_0 = 0$ 出發，預測一系列視差場 ${ d_1, ..., d_N \}$ 。在每次迭代中，利用當前視差估計值對相關體進行索引，生成一組相關特征。這些特征會依次通過兩個卷積層處理，同時當前視差估計值也會經過兩次卷積層處理。隨后將相關特征、視差特征和上下文特征拼接后輸入GRU，由GRU更新隱藏狀態。最終利用更新后的隱藏狀態來預測視差的動態變化。
（1）多個隱藏狀態
原始的Raft算法完全采用固定高分辨率進行更新。這種方法存在一個明顯缺陷：感受野面積會隨著GRU更新次數的增加而緩慢增長，這在處理具有大面積無紋理區域且局部信息匱乏的場景時容易引發問題。為解決這一難題，提出了一種多分辨率更新算子，該算子可同時作用于1/8、1/16和1/32分辨率的特征圖。實驗結果表明，采用這種多分辨率更新算子能顯著提升模型的泛化性能。
如圖2-2所示，GRU通過相互使用對方的隱藏狀態作為輸入進行交叉連接。相關查找和最終的視差更新由分辨率最高的GRU完成。
多級GRU

圖2-2 多級GRU

（2）上采樣
預測的視差場分辨率是輸入圖像的1/4或1/8。為了輸出全分辨率視差圖，采用了與Raft相同的凸上采樣方法。RAFT-Stereo將全分辨率視差值視為其粗分辨率鄰近點3x3網格的凸組合結果，而凸組合權重則由最高分辨率GRU進行預測。
上采樣部分的代碼為

def upsample_flow(self, flow, mask):""" Upsample flow field [H/8, W/8, 2] -> [H, W, 2] using convex combination """N, D, H, W = flow.shapefactor = 2 ** self.args.n_downsamplemask = mask.view(N, 1, 9, factor, factor, H, W)mask = torch.softmax(mask, dim=2)up_flow = F.unfold(factor * flow, [3,3], padding=1)up_flow = up_flow.view(N, D, 9, 1, 1, H, W)up_flow = torch.sum(mask * up_flow, dim=2)up_flow = up_flow.permute(0, 1, 4, 2, 5, 3)return up_flow.reshape(N, D, factor*H, factor*W)

其中的mask就是最高分辨率GRU預測得到的凸組合權重。

2.4 Slow-Fast GRU

將GRU更新到1/8分辨率的隱藏狀態，其浮點運算量大約是更新1/16分辨率隱藏狀態的四倍。為了利用這一特性提升推理速度，作者對RAFT-Stereo模型進行了優化，每當更新1/8分辨率隱藏狀態時，就會同步多次更新1/16和1/32分辨率的隱藏狀態。在KITTI分辨率圖像上進行32次GRU更新后，這一改進使RAFTStereo的運行時間從0.132秒縮短至0.05秒，效率提升達52%。
這種改進使RAFT-Stereo在實時立體視覺方面具有競爭力的性能，而運行的方法快一個數量級，如下表2-1的Slow-Fast行所示。
消融實驗

表2-1 消融實驗

作者發現，通過增加低分辨率GRU的迭代次數并減少高分辨率GRU的迭代次數，在略微降低精度的情況下，RAFT-Stereo的運行時間顯著縮短。Slow-Fast版本的RaftStereo分別將最低、中等和最高分辨率的隱藏狀態更新30次、20次和10次，而Regular版本則將每個隱藏狀態更新32次。無論是Slow-Fast還是Regular版本，都使用相同的模型權重。Slow-Fast-GRU這部分的代碼如下所示

for itr in range(iters):coords1 = coords1.detach()corr = corr_fn(coords1) # index correlation volumeflow = coords1 - coords0with autocast(enabled=self.args.mixed_precision):if self.args.n_gru_layers == 3 and self.args.slow_fast_gru: # Update low-res GRUnet_list = self.update_block(net_list, inp_list, iter32=True, iter16=False, iter08=False, update=False)if self.args.n_gru_layers >= 2 and self.args.slow_fast_gru:# Update low-res GRU and mid-res GRUnet_list = self.update_block(net_list, inp_list, iter32=self.args.n_gru_layers==3, iter16=True, iter08=False, update=False)net_list, up_mask, delta_flow = self.update_block(net_list, inp_list, corr, flow, iter32=self.args.n_gru_layers==3, iter16=self.args.n_gru_layers>=2)

從代碼中不難看出，self.args.slow_fast_gru=True時，計算量反而會增加，并沒有起到提速的效果。這是為啥呢？
使用Slow-Fast-GRU時的做法是將總的迭代次數減少，32->7，但是這樣的效果會變差。但現在模型的做法都是在Slow-Fast-GRU為True的情況下，同時設置迭代次數為32，這其實更慢了。在github issue上面其他人的提問和作者解答可見https://github.com/princeton-vl/RAFT-Stereo/issues/25。

2.5 監督

作者在預測序列 ${ d_1, ..., d_N \}$ 中，通過指數級遞增的權重對預測視差與真實視差之間的 $L 1$ 距離進行監督。給定真實視差 $d_{gt}$ 時，損失函數定義為
$\sum^N_{i=1} \gamma^{N-i} || d_{gt} - d_i ||_1 \tag{2-2}$
其中， $γ=0.9\gamma = 0.9$ 。

3 效果

作者對比了RaftStereo與其他雙目深度估計算法的泛化能力，這些算法均在虛擬場景數據集Sceneflow上進行訓練，并在真實場景數據集上進行測試，可見表3-1。表中的數值表示誤差，即誤差超過指定閾值的像素百分比。作者采用標準評估閾值是KITTI為3像素，Middlebury為2像素，ETH3D為1像素。
泛化能力效果對比