拋棄編碼器-解碼器架構,用擴散模型做邊緣檢測效果更好,國防科大提出DiffusionEdge
現有的深度邊緣檢測網絡通常基于包含了上下采樣模塊的編碼器 - 解碼器架構,以更好的提取多層次的特征,但這也限制了網絡輸出既準又細的邊緣檢測結果。
針對這一問題,一篇 AAAI 2024 上的論文給出了新的解決方案。
- 論文題目:DiffusionEdge: Diffusion Probabilistic Model for Crisp Edge Detection
- 作者:葉云帆(國防科技大學),徐凱(國防科技大學),黃雨行(國防科技大學),易任嬌(國防科技大學),蔡志平(國防科技大學)
- 論文鏈接: https://arxiv.org/abs/2401.02032
- 開源代碼: https://github.com/GuHuangAI/DiffusionEdge
國防科技大學 iGRAPE Lab 提出了首個用于二維邊緣檢測任務的擴散概率模型方法,通過學習迭代的去噪過程獲得邊緣結果圖,去噪過程如圖 1 所示。為了在保留最終性能的同時減少計算資源的消耗,該方法在隱空間中訓練網絡,并引入了不確定性蒸餾模塊以更好的優化。同時,本方法還采用解耦架構來加速去噪過程,并提出了相應的自適應傅立葉濾波器來調整特征。基于這些設計,本方法得以用有限的資源進行穩定的訓練,用更少的增強策略來預測清晰準確的邊緣圖。在四個公共基準數據集上的大量實驗表明,本文提出的方法在準確度和粗細度上全面超越了其他方法。
圖 1 基于擴散概率模型的邊緣檢測過程與優勢示例
本文的創新點包括:
1、提出了領域內第一個針對邊緣檢測任務的擴散模型 DiffusionEdge,無需任何后處理即可預測出更細更準確的邊緣圖。
2、針對擴散模型應用時的難點,設計了多種技術以確保方法能在隱空間中穩定學習,同時保留像素級的不確定性先驗知識并自適應地過濾傅立葉空間中的隱特征。
3、在四個邊緣檢測公共基準數據集上開展的大量對比實驗展示了 DiffusionEdge 在準確度和細度方面均具有卓越的性能優勢。
相關工作
基于深度學習的方法通常采用包含上下采樣的編解碼結構集成多層特征 [1-2],或是整合多個標注的不確定性信息以提升邊緣檢測的準確度 [3]。然而,天然受限于這樣的結構,其生成的邊緣結果圖對于下游任務來說太過粗厚而嚴重依賴后處理的問題仍然亟待解決。盡管許多工作已經在損失函數 [4-5] 和標簽修正策略 [6] 方面做出了探索以使網絡能輸出更細的邊緣,但本文認為該領域仍然需要一個可以不借助任何額外模塊,就能直接滿足準確度和細度的邊緣檢測器,而無需任何后處理步驟。
擴散模型是一類基于馬爾可夫鏈的生成模型,通過學習去噪過程逐漸恢復目標數據樣本。擴散模型在計算機視覺、自然語言處理和音頻生成等領域都表現出了卓越的性能。不僅如此,通過將圖像或是其他模態的輸入作為額外條件時,其在感知任務中也表現出了巨大的潛力,例如圖像分割 [7]、目標檢測 [8] 和姿態估計 [9] 等。
方法描述
本文所提出的 DiffusionEdge 方法總體框架如圖 2 所示。受以往工作的啟發,該方法在隱空間中訓練具有解耦結構的擴散模型,并將圖像作為額外的條件線索輸入。該方法引入了自適應傅里葉濾波器來進行頻率解析,且為了保留來自多個標注者的像素級不確定性信息并減少對計算資源的要求,還以蒸餾的方式直接使用交叉熵損失優化隱空間。
圖 2 DiffusionEdge 的整體結構示意
針對目前的擴散模型受到采樣步數太多,推理時間太長等問題的困擾,本方法受 DDM [10] 的啟發,同樣使用解耦的擴散模型架構來加速采樣推理過程。其中,解耦的前向擴散過程由顯式的轉移概率和標準 Wiener 過程的組合來控制:
其中和
分別代表初始邊緣和噪聲邊緣,
指的是反向邊緣梯度的顯式轉換函數。與 DDM 類似,本文方法默認使用常量函數
,則其相應的逆過程可以被表示為:
其中。為了訓練解耦的擴散模型,該方法需要同時監督數據和噪聲分量,因此,訓練目標可以被參數化為:
其中是去噪網絡中的參數。由于擴散模型如果要在原始圖像空間中訓練的話會占用太多的計算成本,因此參考 [11] 的思路,本文提出的方法將訓練過程轉移到了具有 4 倍下采樣空間大小的隱空間中。
如圖 2 所示,該方法首先訓練了一對自編碼器和解碼器的網絡,該編碼器將邊緣標注壓縮為一個隱變量,而解碼器則用于從這個隱變量中恢復出原來的邊緣標注。如此一來,在訓練基于 U-Net 結構的去噪網絡階段,該方法便固定這一對自編碼和解碼器網絡的權重,并在隱空間中訓練去噪過程,這樣可以大幅減少網絡對計算資源的消耗,同時維持不錯的性能。
為了提升網絡最后的性能,本文提出的方法在解耦操作中引入了一種可以自適應地過濾掉不同頻率特征的模塊。如圖 2 左下角所示,該方法將自適應的快速傅立葉變換濾波器(Adaptive FFT-filter)集成到了去噪 Unet 網絡中,位置在解耦操作之前,以在頻域中自適應過濾并分離出邊緣圖和噪聲分量。具體來說,給定編碼器特征 ,該方法首先沿空間維度執行二維的傅里葉變換(FFT),并將變換后的特征表示為
。緊接著,為了訓練這個自適應頻譜濾波模塊,構造了一個可學習的權重圖
并將其W乘以Fc。頻譜濾波器可以全局地調整特定頻率,并且學習到的權重可以適應不同數據集目標分布的不同頻率情況。通過自適應地濾除無用成分,本方法通過快速傅里葉逆變換(IFFT)操作將特征從頻域映射回空間域。最后,通過額外引入了來自的殘差連接,避免完全過濾掉了所有有用的信息。上述過程可以由以下公式來描述:
其中是輸出特征,o表示哈達瑪積(Hadamard Product)。
由于邊緣和非邊緣像素的數量高度不平衡(大多數像素都是非邊緣的背景),參考以往的工作,我們同樣引入不確定性感知的損失函數進行訓練。具體來說,將作為第i個像素的真值邊緣概率,對于第j個邊緣圖中的第i個像素,其值為,則不確定性感知 WCE 損失的計算方式如下:
其中,其中
是決定真值標注中不確定的邊緣像素的閾值,如果像素值大于 0 而小于此閾值,則此類模糊的,置信度不夠高的像素樣本將在后續的優化過程中被忽略(損失函數為 0)。
和
分別表示真值標注邊緣圖中邊緣和非邊緣像素的數量。是用來平衡
和
的權重(設為 1.1)。因此,每個邊緣圖的最終損失函數計算方式為
。
在優化過程中忽略模糊的低置信度的像素可以避免網絡混亂,使訓練過程更穩定的收斂,并提高模型的表現。然而,將二元交叉熵損失直接應用于在數值和空間上均未對齊的隱空間中是幾乎不可能的。尤其是,不確定性感知的交叉熵損失中利用閾值(一般從 0 到 1)來判斷像素是否為邊緣,這是從圖像空間上定義的,而隱變量遵循正態分布并且具有完全不同的范圍和實際意義。此外,像素級的不確定性很難與不同大小的編碼和下采樣的隱特征保持一致,二者是很難直接兼容的。因此,將交叉熵損失直接應用于優化隱變量不可避免地會導致不正確的不確定性感知。
另一方面,人們可以選擇將隱變量解碼回圖像級別,從而可以直接使用不確定性感知交叉熵損失來直接監督預測出的邊緣結果圖。不幸的是,這種實現讓反向傳播的參數梯度通過了冗余的自編碼器網絡,使得梯度很難得到有效的傳遞。此外,自編碼器網絡中額外的梯度計算會帶來巨大的 GPU 內存消耗成本,這違背了本方法設計一個實用的邊緣檢測器的初衷,很難推廣到實際應用中。因此,本方法提出了不確定性蒸餾損失,可以直接優化隱空間上的梯度,具體來說,設重構出的隱變量為,自編碼器網絡的解碼器是D,解碼出的邊緣結果是eD,本方法考慮基于鏈式法則直接計算不確定性感知的二元交叉熵損失
的梯度,具體計算方式為:
為了消除自編碼器網絡的負面影響,本方法直接跳過了自編碼器傳遞梯度并將梯度
的計算方式修改調整為:
這樣的實現方式大大降低了計算成本,并允許利用不確定性感知的損失函數直接在隱變量上優化。如此一來,再結合一個隨步數t自適應變化的時變損失權重,本方法最終訓練優化目標可以被表示為:
實驗結果
本方法在四個領域內被廣泛使用的邊緣檢測公共標準數據集上進行了實驗:BSDS、NYUDv2、Multicue 和 BIPED。由于邊緣檢測數據標注比較難,標注數據量都比較少,以往的方法通常會使用各種策略來增強數據集。例如說,BSDS 中的圖像通過水平翻轉 (2×)、縮放 (3×) 和旋轉 (16×) 進行增強,能夠生成比原始版本擴大了 96 倍的訓練集。以往的方法在其他數據集上使用的通用增強策略總結在了表格 1 中,其中 F 代表水平翻轉,S 代表縮放,R 代表旋轉,C 代表裁剪,G 代表伽馬校正。不同的是,本方法僅需要使用隨機裁剪的 320320 的圖像塊來訓練所有數據。在 BSDS 數據集中,本方法則僅僅采用隨機的翻轉和縮放,其定量對比結果展示在了表 2 中。在 NYUDv2、Mu lticue 和 BIPED 數據集中,該方法僅需采用隨機翻轉訓練。在使用更少增強策略的情況下,本方法在各個數據集,各個指標上的表現都優于之前的方法。通過觀察圖 3-5 的預測結果可以看出,DiffusionEdge 能夠學習并預測出跟 gt 分布幾乎一樣的邊緣檢測結果圖來,預測結果準確而清晰的優勢對于有精細化需求的下游任務來說非常重要,也展示了其能直接應用于后續任務的巨大潛力。
表 1 以往方法在四個邊緣檢測數據集上使用的增強策略
表 2 不同方法在 BSDS 數據集上的定量對比
圖 3 不同方法在 BSDS 數據集上的定性對比
圖 4 不同方法在 NYUDv2 數據集上的定性對比
圖 5 不同方法在 BIPED 數據集上的定性對比