一種基于Transform的神經(jīng)網(wǎng)絡(luò)架構(gòu)

摘要

大家好，今天為大家?guī)?lái)的文章 A Transformer Architecture for Optical Flow 我們介紹了光流估計(jì)網(wǎng)絡(luò)，稱為FlowFormer，一種基于Transform的神經(jīng)網(wǎng)絡(luò)架構(gòu)，用于學(xué)習(xí)光流。

FlowFormer化由圖像對(duì)構(gòu)建的4D cost volume，將成本編碼到一個(gè)新的潛在空間中具有交替組轉(zhuǎn)換器(AGT)層的成本內(nèi)存中，并通過(guò)一個(gè)帶有動(dòng)態(tài)位置成本查詢的循環(huán)Transform解碼器對(duì)位置 cost queries進(jìn)行解碼。

在sinintel基準(zhǔn)測(cè)試中，F(xiàn)lowFormer的平均終點(diǎn)誤差(AEPE)分別為1.159和2.088，比已發(fā)布的最佳結(jié)果誤差分別降低了16.5%和15.5%

主要工作與貢獻(xiàn)

我們的貢獻(xiàn)可以概括為四個(gè)方面。

1)我們提出了一種新的基于Transform的神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu)FlowFormer，用于光流量估計(jì)，它實(shí)現(xiàn)了最先進(jìn)的流量估計(jì)性能。

2)設(shè)計(jì)了一種新穎的cost volume編碼器，有效地將成本信息聚合為緊湊的潛在cost tokens。

3)我們提出了一種循環(huán)成本解碼器，該解碼器通過(guò)動(dòng)態(tài)位置成本查詢循環(huán)解碼成本特征，迭代細(xì)化估計(jì)光流。

4)據(jù)我們所知，我們第一次驗(yàn)證imagenet預(yù)先訓(xùn)練的傳輸

方法

光流估計(jì)任務(wù)要求輸出逐像素位移場(chǎng)f: R2→R2，將源圖像的每個(gè)二維位置x∈R2 Is映射到目標(biāo)圖像It對(duì)應(yīng)的二維位置p = x+f (x)。為了充分利用現(xiàn)有的視覺(jué)Transform體系結(jié)構(gòu)，以及之前基于cnn的光流估計(jì)方法廣泛使用的四維成本體積，我們提出了一種基于Transform的結(jié)構(gòu)FlowFormer，它對(duì)四維成本體積進(jìn)行編碼和解碼，以實(shí)現(xiàn)精確的光流估計(jì)。

在圖1中，我們展示了FlowFormer的總體架構(gòu)，它處理來(lái)自siamese特征的4D成本卷，包含兩個(gè)主要組件: 1)cost volume編碼器，將4D成本卷編碼到潛在空間中形成cost memory; 2)成本記憶解碼器，用于基于編碼的成本記憶和上下文特征預(yù)測(cè)每像素位移場(chǎng)。



圖1 FlowFormer的體系結(jié)構(gòu)。

FlowFormer通過(guò)三個(gè)步驟估算光流量:1)根據(jù)圖像特征構(gòu)建4Dcost volume。2)成本卷編碼器，將成本卷編碼到成本內(nèi)存中。3)循環(huán)Transform解碼器，將具有源圖像上下文特征的代價(jià)內(nèi)存解碼為流

1.1 構(gòu)建 4D Cost Volume

骨干網(wǎng)用于從輸入HI × WI × 3rgb圖像中提取H ×W ×Df特征圖，通常我們?cè)O(shè)置(H, W) = (HI /8, WI /8)。在提取出源圖像和目標(biāo)圖像的特征圖后，通過(guò)計(jì)算源圖像和目標(biāo)圖像之間所有像素對(duì)的點(diǎn)積相似度，構(gòu)造H × W × H × W 4D代價(jià)體。

1.2 Cost Volume Encoder 為了估計(jì)光流，需要基于4D代價(jià)體中編碼的源目標(biāo)視覺(jué)相似性來(lái)識(shí)別源像素在目標(biāo)圖像中的對(duì)應(yīng)位置。構(gòu)建的4D成本體積可以被視為一系列大小為H × W的2D成本圖，每個(gè)成本圖度量單個(gè)源像素和所有目標(biāo)像素之間的視覺(jué)相似性。

我們將源像素x的代價(jià)映射表示為Mx∈RH×W。

在這樣的成本圖中找到對(duì)應(yīng)的位置通常是具有挑戰(zhàn)性的，因?yàn)樵趦蓚€(gè)圖像中可能存在重復(fù)的模式和非歧視性區(qū)域。當(dāng)只考慮來(lái)自地圖局部窗口的成本時(shí)，任務(wù)變得更加具有挑戰(zhàn)性，就像以前基于cnn的光流估計(jì)方法所做的那樣。即使是估計(jì)單個(gè)源像素的精確位移，考慮其上下文源像素的代價(jià)圖也是有益的。

為了解決這一難題，我們提出了一種基于Transform的成本體積編碼器，該編碼器將整個(gè)成本體積編碼到成本內(nèi)存中。我們的成本卷編碼器包括三個(gè)步驟:1)成本映射補(bǔ)丁化，2)成本補(bǔ)丁嵌入，3)成本內(nèi)存編碼。我們?cè)敿?xì)闡述這三個(gè)步驟如下。

1.2 地圖patchification成本

我們根據(jù)已有的視覺(jué)轉(zhuǎn)換器，對(duì)每個(gè)源像素x的代價(jià)映射Mx∈RH×W進(jìn)行跨卷積拼接，得到代價(jià)補(bǔ)丁嵌入序列。具體來(lái)說(shuō)，給定一個(gè)H ×Wcost圖，我們首先在它的右側(cè)和底部填充0，使其寬度和高度為8的倍數(shù)。然后，填充的代價(jià)映射通過(guò)三個(gè)stride-2卷積的堆棧和ReLU轉(zhuǎn)換為特征映射Fx∈R?H/8?×?W/8?×Dp。特征圖中的每個(gè)特征代表輸入成本圖中的一個(gè)8 × 8補(bǔ)丁。三種卷積的輸出通道均為Dp / 4, Dp / 2 Dp。

1.3基于潛在摘要的補(bǔ)丁特征標(biāo)記

盡管對(duì)每個(gè)源像素進(jìn)行修補(bǔ)會(huì)得到一系列cost patch特征向量，但這種patch特征的數(shù)量仍然很大，影響了信息在不同源像素間傳播的效率。實(shí)際上，成本圖是高度冗余的，因?yàn)橹挥猩贁?shù)高成本是最有信息的。為了獲得更緊湊的代價(jià)特征，我們進(jìn)一步通過(guò)K個(gè)潛碼字C∈RK×D總結(jié)了每個(gè)源像素x的patch特征Fx。

具體來(lái)說(shuō)，latent codewords query每個(gè)源像素的cost補(bǔ)丁特征，通過(guò)點(diǎn)積注意機(jī)制將每個(gè)成本圖進(jìn)一步總結(jié)為K個(gè)D維的潛在向量。潛碼字C∈RK×D隨機(jī)初始化，通過(guò)反向傳播進(jìn)行更新，并在所有源像素之間共享。歸納Fx的潛在表示Tx得到為



在投影成本-補(bǔ)丁特征Fx以獲得鍵Kx和值Vx之前，將補(bǔ)丁特征與位置嵌入序列PE∈R?H/8?×?W/8?×Dp進(jìn)行拼接。給定一個(gè)2D位置p，我們將其編碼為一個(gè)長(zhǎng)度為Dp的位置嵌入，跟隨COTR[27]。最后，通過(guò)對(duì)查詢、鍵和值進(jìn)行多頭點(diǎn)積注意，可以將源像素x的代價(jià)圖匯總為K個(gè)潛在表示Tx∈RK×D。

關(guān)注潛在cost空間。上述兩個(gè)階段將原始的4D成本體積轉(zhuǎn)化為潛在的、緊湊的4D成本體積t。但是，直接對(duì)4D體積中的所有向量應(yīng)用自注意仍然過(guò)于昂貴，因?yàn)橛?jì)算成本隨tokens數(shù)量的增加呈二次增長(zhǎng)。 如圖2所示，我們提出了一種交替分組轉(zhuǎn)換層(AGT)，該層以兩種相互正交的方式對(duì)標(biāo)記進(jìn)行分組，并在兩組中交替應(yīng)用注意，減少了注意成本，同時(shí)仍然能夠在所有標(biāo)記之間傳播信息。



圖2 Alternate-Group Transform層。

交替組轉(zhuǎn)換層(AGT)將T中的token交替分組為包含K token的H × W組(Tx)和包含H × W token的K組(Ti)，并分別通過(guò)self-attention和ss self-attention[8]對(duì)組內(nèi)的token進(jìn)行編碼。 對(duì)每個(gè)源像素進(jìn)行第一次分組，即每個(gè)Tx∈RK×D組成一個(gè)組，在每個(gè)組內(nèi)進(jìn)行自我注意。



在每個(gè)源像素x的所有K個(gè)潛在令牌之間進(jìn)行自注意后，更新的Tx通過(guò)前饋網(wǎng)絡(luò)(FFN)進(jìn)一步變換，然后重新組織，形成更新的4D成本體積t。自注意子層和FFN子層均采用Transform殘差連接和層歸一化的共同設(shè)計(jì)。這種自我注意操作在每個(gè)成本映射內(nèi)傳播信息，我們將其命名為內(nèi)部成本映射的自我注意。

第二種方法根據(jù)K種不同的潛在表示，將所有潛在代價(jià)令牌T∈RH×W ×K×D分成K組。因此，每一組都有D維度的(H × W)標(biāo)記，用于通過(guò)雙胞胎[8]中提出的空間可分離的自我注意(SS-SelfAttention)在空間域中傳播信息，

上述自我注意操作的參數(shù)在不同組之間共享，并按順序進(jìn)行操作，形成提出的交替組注意層。通過(guò)多次疊加交替組Transform層，潛在成本令牌可以有效地跨源像素和跨潛在表示交換信息，以更好地編碼4D成本體積。通過(guò)這種方式，我們的成本體積編碼器將H × W × H × W 4D成本體積轉(zhuǎn)換為H × W × K長(zhǎng)度為d的潛在標(biāo)記。我們將最終的H × W × K標(biāo)記稱為成本存儲(chǔ)器，用于光學(xué)解碼。

1.4用于流量估計(jì)的成本記憶解碼器

考慮到成本體積編碼器編碼的cost記憶，我們提出了一種用于光流預(yù)測(cè)的cost記憶解碼器。由于輸入圖像的原始分辨率為HI × WI，我們?cè)贖 × W分辨率下估計(jì)光流，然后使用可學(xué)習(xí)的凸上采樣器[46]對(duì)預(yù)測(cè)的流進(jìn)行上采樣到原始分辨率。然而，與以往的視覺(jué)Transform尋找抽象的語(yǔ)義特征不同，光流估計(jì)需要從記憶中恢復(fù)密集的對(duì)應(yīng)。

受RAFT[46]的啟發(fā)，我們建議使用成本查詢從cost內(nèi)存中檢索成本特征，并使用循環(huán)注意解碼器層迭代優(yōu)化流量預(yù)測(cè)。

1.5Cost memory aggregation:

為了預(yù)測(cè)H × W源像素的流，我們生成了一個(gè)(H × W)成本查詢序列，每個(gè)(H × W)成本查詢通過(guò)對(duì)成本內(nèi)存的共同關(guān)注負(fù)責(zé)估計(jì)單個(gè)源像素的流。為了為源像素x生成成本查詢Qx，我們首先計(jì)算其在目標(biāo)圖像中的對(duì)應(yīng)位置，給定其當(dāng)前估計(jì)流量f (x)，即p = x + f (x)。

然后，通過(guò)裁剪成本映射Mx上以p為中心的9×9本地窗口內(nèi)的成本，我們檢索一個(gè)局部9×9成本映射補(bǔ)丁qx = Crop9×9(Mx, p)。然后根據(jù)局部成本Qx編碼的特征FFN(Qx)和p的位置嵌入PE(p)構(gòu)造成本查詢Qx，通過(guò)交叉注意聚合源像素x的成本記憶Tx信息。

循環(huán)光流預(yù)測(cè)。我們的成本解碼器迭代回歸流量殘差?f (x)，將每個(gè)源像素x的流量細(xì)化為f (x)←f (x) +?f (x)。我們采用ConvGRU模塊，并遵循與GMARAFT[25]類似的設(shè)計(jì)進(jìn)行流量細(xì)化。

然而，我們的循環(huán)模塊的關(guān)鍵區(qū)別是使用成本查詢來(lái)自適應(yīng)地聚合來(lái)自成本內(nèi)存的信息，以實(shí)現(xiàn)更精確的流量估計(jì)。具體來(lái)說(shuō)，在每次迭代中，ConvGRU單元將檢索到的成本特征和成本映射補(bǔ)丁Concat(cx, qx)、來(lái)自上下文網(wǎng)絡(luò)的源圖像上下文特征tx和當(dāng)前估計(jì)的流量f的拼接作為輸入，并輸出預(yù)測(cè)的流量殘差如下

實(shí)驗(yàn)結(jié)果

?

圖4 Sintel[3]和KITTI[14]數(shù)據(jù)集的實(shí)驗(yàn)。

' A '表示自動(dòng)流數(shù)據(jù)集。“C + T”表示只在FlyingChairs和FlyingThings數(shù)據(jù)集上進(jìn)行訓(xùn)練。' + S + K + H '表示對(duì)Sintel、KITTI和HD1K訓(xùn)練集的組合進(jìn)行微調(diào)。*表示這些方法使用warm-start策略[46]，該策略依賴于視頻中先前的圖像幀。?是通過(guò)補(bǔ)充闡述的瓦技術(shù)估計(jì)的。我們的FlowFormer實(shí)現(xiàn)了最佳的泛化性能(C+T)，并在Sintel基準(zhǔn)測(cè)試(C+T+S+K+H)中排名第一。



圖5 在steint檢驗(yàn)集上的定性比較。

FlowFormer大大減少了對(duì)象邊界(紅色箭頭所指)周圍的流動(dòng)泄漏，并使細(xì)節(jié)(藍(lán)色箭頭所指)更加清晰。


圖5 在steint檢驗(yàn)集上的定性比較。

FlowFormer大大減少了對(duì)象邊界(紅色箭頭所指)周圍的流動(dòng)泄漏，并使細(xì)節(jié)(藍(lán)色箭頭所指)更加清晰。




審核編輯：劉清