日前, 美圖影像研究院 (MT Lab) 聯(lián)合北京交通大學(xué)提出了超高分辨率摳圖方法 MEMatte (MemoryEfficient), 一個(gè)顯存友好的自然圖像摳圖框架用于在顯存受限的場(chǎng)景下 (如商用顯卡與邊緣設(shè)備) 進(jìn)行高清圖像的精細(xì)摳圖。文章已被人工智能頂會(huì) AAAI 2025 接收。

隨著攝影設(shè)備的性能升級(jí), 我們?cè)谌粘２杉降膱D像分辨率已經(jīng)能達(dá)到 4K 乃至 8K, 這種高分辨率圖像有著更好的視覺質(zhì)感與更豐富的細(xì)節(jié)內(nèi)容。而自然圖像摳圖方法相比于一般的分割方法更善于區(qū)分細(xì)節(jié)充實(shí)的場(chǎng)景, 因此適用于更復(fù)雜的下游任務(wù)諸如圖像與視頻的后期編輯。

當(dāng)前 Transformer (ViT) 為自然圖像摳圖方法帶來了顯著的性能提升, 這種提升來自于全局注意力模塊對(duì)圖像內(nèi)容的長(zhǎng)程建模能力。然而, 基于 ViT 的摳圖方法在計(jì)算資源受限的場(chǎng)景下無法處理高分辨率圖像, 因?yàn)?ViT 內(nèi)部的全局注意力會(huì)產(chǎn)生次方倍的計(jì)算開銷。此外, 摳圖任務(wù)的特性要求輸入能同時(shí)保留清晰圖像中的細(xì)節(jié)信息與完整圖像中的語義信息, 因此無法采用降采樣或圖像切片等手段來降低模型輸入的圖像尺寸。這些限制條件要求我們?cè)诰S持圖像輸入大小的前提下降低全局注意力對(duì)計(jì)算資源的消耗。

過往的解決方案包括令牌剪枝與令牌融合, 這些方法都需要在模型的不同層以固定比例逐漸丟棄冗余令牌。但是這類方法存在兩種缺陷:

(1) 丟棄令牌會(huì)導(dǎo)致?lián)笀D效果顯著變差, 因?yàn)楸粊G棄的令牌中存在著需要被保留的細(xì)節(jié)信息;

(2) 預(yù)設(shè)固定的丟棄比例無法適應(yīng)不同的輸入, 對(duì)于復(fù)雜樣例丟棄過多信息會(huì)導(dǎo)致精度下降, 而對(duì)于簡(jiǎn)單樣例則依然存在計(jì)算冗余。

為了解決上述問題, 研究團(tuán)隊(duì)提出的 MEMatte 采用了雙分支的令牌路由設(shè)計(jì)。該設(shè)計(jì)的核心思想是將包含語義信息的令牌送入全局注意力模塊進(jìn)行處理, 并且將其余令牌送入精心設(shè)計(jì)的輕量化令牌提煉模塊 (Light-weight Token Refinement Module, LTRM) 進(jìn)行處理。令牌的分流是通過動(dòng)態(tài)路由機(jī)制 (Batch-constrained Adaptive Token Routing, BATR) 進(jìn)行篩選, 在推理時(shí)無需預(yù)設(shè)固定的比例, 因此能夠輸入自適應(yīng)地進(jìn)行調(diào)整。這一設(shè)計(jì)顯著降低了模型的計(jì)算開銷, 并在 Nvidia GeForce 2080Ti 商用 GPU 上實(shí)現(xiàn)了 4K 分辨率圖像摳圖。此外, 研究團(tuán)隊(duì)還開源了超高分辨率自然圖像摳圖數(shù)據(jù)集 UHR-395 (Ultra High Resolution dataset), 用于高分辨率模型的訓(xùn)練與評(píng)估。

圖 1:MEMatte 方法與已有方法在不同分辨率的顯存消耗對(duì)比對(duì)比

圖 2:UHR-395 數(shù)據(jù)集與其它開源摳圖數(shù)據(jù)集的輸入下 UHR-395 有著最高的平均分辨率。

方法

在推理過程中,MEMatte 的 ViT 編碼器內(nèi)部的每一個(gè)全局注意力模塊前都存在一個(gè)路由器 (Router)。這個(gè)路由器通過局部-全局策略為全局注意力進(jìn)行令牌重要性評(píng)估。該策略對(duì)每層的圖像輸入進(jìn)行重新映射以及通道拆分后的全局池化操作來兼顧局部與全局信息。此后通過 LogSoftmax (LS) 操作形成令牌的分流結(jié)果。

Router 將分流到全局注意力之外的令牌送入 LTRM 進(jìn)行處理。LTRM 由數(shù)個(gè)輕量化組建構(gòu)成, 其中包括兩組映射層。此外, 深度卷積 (Depth-Wise Convolution, DWC) 層處理局部空間信息, 高效通道注意力層 (Efficient Channel Attention, ECA) 處理全局特征信息。

在訓(xùn)練 MEMatte 時(shí), 研究團(tuán)隊(duì)采用 BATR 的訓(xùn)練機(jī)制, 搭配蒸餾學(xué)習(xí)策略與常規(guī)摳圖訓(xùn)練損失, 實(shí)現(xiàn)快速且平穩(wěn)的模型訓(xùn)練, 在訓(xùn)練過程中, 假定從數(shù)據(jù)分布中均勻地采樣, 則每次的壓縮比例應(yīng)當(dāng)為一個(gè)定值?；谶@項(xiàng)假設(shè), BATR 為模型預(yù)設(shè)了分流比例

來約束模型預(yù)測(cè)的批次平均分流比例, 定義批次平均分流比例

為如下形式:

通過約束模型的預(yù)測(cè)值與預(yù)設(shè)值的差異, 訓(xùn)練出足夠有效的 Router 以自適應(yīng)的處理各種圖像。

圖 3:MEMatte 的訓(xùn)練與推理流程

試驗(yàn)結(jié)果

實(shí)驗(yàn)結(jié)果表明,MEMatte 在現(xiàn)有的基準(zhǔn)以及 UHR-395 的高分辨率測(cè)試集上都領(lǐng)先于過去的方法, 并在域外測(cè)試集上展現(xiàn)了良好的泛化性能。此外, 在摳圖測(cè)試集 Composition-1K 上,MEMatte 相比于其 baseline 節(jié)省了約 88% 的顯存開銷, 并降低了約 50% 的推理時(shí)間。

圖 4:MEMatte 與過往方法在通用摳圖基準(zhǔn)上的比較

圖 5:MEMatte 與過往方法在 UHR-395 高分辨率測(cè)試集上的比較

圖 6:在高分辨率圖像上的摳圖效果。

研究圖案對(duì) Router 分流結(jié)果的可視化結(jié)果展示了 MEMatte 在不同的全局注意力層對(duì)令牌的選擇傾向。MeMatte 使中間的自注意力層主要關(guān)注物體邊緣與細(xì)區(qū)域的復(fù)雜區(qū)域, 同時(shí)避免了令牌通過 ViT 淺層的自注意力操作, 這大幅提高了輸入自注意力層的令牌稀疏性。

圖 7:對(duì) Router 分流結(jié)果的可視化。其中灰色區(qū)域是分流向 LTRM 的令牌。

超高分辨率自然圖像摳圖數(shù)據(jù)集

目前公開的摳圖數(shù)據(jù)集難以作為高分辨率摳圖的基準(zhǔn), 因?yàn)樗鼈儫o法同時(shí)兼顧類別多樣性和高分辨率。為了解決這一問題, 研究團(tuán)隊(duì)提出了一個(gè)超高分辨率自然摳圖數(shù)據(jù)集 UHR-395。該數(shù)據(jù)集包含 11 個(gè)類別, 共 395 個(gè)前景物體, 劃分為訓(xùn)練集 355 張以及測(cè)試集 40 張, 平均分辨率高達(dá) 4872×6017。

圖 8:前景類別分布。

圖 9:圖像分辨率分布。

為確保標(biāo)注數(shù)據(jù)的高質(zhì)量, 每個(gè)前景物體的標(biāo)注結(jié)果都通過多輪嚴(yán)格的審核機(jī)制, 并由摳圖領(lǐng)域?qū)＜疫M(jìn)行最后確認(rèn), 不合格標(biāo)注結(jié)果均被返回進(jìn)行重新修正, 最終獲得高質(zhì)量的自然圖像摳圖標(biāo)注數(shù)據(jù)。

圖 10:標(biāo)注結(jié)果展示。

結(jié)語

超高分辨率的自然圖像摳圖方法對(duì)實(shí)現(xiàn)復(fù)雜圖像編輯與視頻后期處理等下游任務(wù)而言至關(guān)重要, 在本研究中, 來自美圖影像研究院 (MT Lab) 和北京交通大學(xué)的研究人員共同提出了 MEMatte 和超高分辨率數(shù)據(jù)集 UHR-395。與現(xiàn)有的摳圖方法相比, 采用 MEMatte 能夠顯著降低推理過程中的內(nèi)存使用量和延遲, 在高分辨率和真實(shí)世界數(shù)據(jù)集上實(shí)現(xiàn)了最先進(jìn)的性能, 突出了令牌路由機(jī)制 BATR 和 輕量化令牌細(xì)化模塊 LTRM 的穩(wěn)健性, 為美圖的圖像與視頻處理能力提供高性能且靈活的技術(shù)支撐, 為用戶帶來高效優(yōu)質(zhì)使用體驗(yàn)。此外, MEMatte 也有望促進(jìn)高分辨率摳圖的進(jìn)一步研究。