近日網(wǎng)易互娛 AI Lab 獲得第七屆 NTIRE HDR 比賽的全部?jì)蓚€(gè)賽道的冠軍。網(wǎng)易互娛 AI Lab 一直致力于利用 AI 提升美術(shù)生產(chǎn)效率,助力游戲貼圖資源自動(dòng)升級,目前相關(guān)技術(shù)已應用于貼圖、UI 等游戲資源的精度和細節的提升,為多個(gè)游戲提供技術(shù)支持。本文將詳細解讀他們的雙冠比賽方案。
近日,圖像修復領(lǐng)域最具影響力的國際頂級賽事——New Trends in Image Restoration and Enhancement(NTIRE)結果出爐,網(wǎng)易互娛 AI Lab 包攬了高動(dòng)態(tài)范圍成像(HDR)任務(wù)全部 2 項冠軍。NTIRE 比賽每年舉辦一次,目前已是第七屆,主要聚焦圖像修復和增強技術(shù),代表相關(guān)領(lǐng)域的趨勢和發(fā)展,吸引了眾多來(lái)自工業(yè)界、學(xué)術(shù)界的關(guān)注者和參賽者,有著(zhù)非常大的影響力。今年 NTIRE 比賽在計算機視覺(jué)領(lǐng)域的頂級會(huì )議 CVPR 2022(Computer Vision and Pattern Recognition)上舉辦。
高動(dòng)態(tài)范圍成像(HDR)任務(wù)的賽道 1 和賽道 2 分別有 197 個(gè)隊伍、168 個(gè)隊伍報名參加,吸引了包括騰訊、頭條、曠視、螞蟻、快手在內的工業(yè)界隊伍,以及清華大學(xué)、中科院、中國科學(xué)技術(shù)大學(xué)、愛(ài)丁堡大學(xué)、帝國理工等國內外高校。網(wǎng)易互娛 AI Lab 從眾多的強隊中脫穎而出,斬獲該任務(wù)的全部 2 項冠軍。這是網(wǎng)易互娛 AI Lab 奪得多項國際冠軍后,再次登頂國際 AI 競賽,展現了網(wǎng)易互娛 AI Lab 在人工智能領(lǐng)域的綜合實(shí)力。
任務(wù)描述
消費級的單傳感器相機在拍攝照明情況復雜的場(chǎng)景時(shí),難以用一種曝光參數拍攝出曝光正常的照片。由于成像傳感器固有的局限性,譬如在高亮度區域因為過(guò)曝讓畫(huà)面呈現白色,從而導致這些區域細節的丟失。針對這個(gè)問(wèn)題的解決方法是采用一個(gè)擴展的高動(dòng)態(tài)范圍(HDR)記錄圖片,具體做法是在拍攝時(shí)分辨拍攝多張不同曝光參數的低動(dòng)態(tài)范圍(LDR)圖片,然后通過(guò)算法把多張圖片融合成一張 HDR 圖片。
此次比賽的目標是探索高效的 HDR 模型和技術(shù)方案,以達到實(shí)用化的使用需求??傆媰蓚€(gè)賽道:(1)保真度賽道:在限定模型計算量(小于 200G GMACs)的基礎上,盡可能取得更高的保真度,評價(jià)指標是 PSNR-μ;(2)低復雜度賽道:在超過(guò)基線(xiàn)模型指標(PSNR-μ與 PSNR)的基礎上,盡可能取得更低的計算量和更少的耗時(shí),評價(jià)指標是 GMACs。
數據集介紹
本次比賽用的數據集包含 1500 個(gè)訓練樣本、60 個(gè)驗證樣本以及 201 個(gè)測試樣本,每個(gè)樣本包括三張 LDR 圖片輸入,分別對應短、中、長(cháng)三種曝光水平,以及一個(gè)和中曝光對齊的 HDR 圖片,數據集是由 Froehlich 等人收集的,他們捕捉了各種各樣的具有挑戰性場(chǎng)景的 HDR 視頻。之前基于深度學(xué)習的 HDR 模型取得了不錯的效果,譬如 AHDRNet、ADNet 等,但缺點(diǎn)是計算量非常大,以官方提供的基線(xiàn)方法 AHDRNet 為例,計算量在 3000GMACs 左右。因此本次比賽的目的是尋求高效的多幀 HDR 重建方法。
方法概述
由于任務(wù)的兩個(gè)賽道均要求訓練高效的 HDR 模型,網(wǎng)易互娛 AI Lab 憑借以往對 low-level 視覺(jué)任務(wù)和輕量化網(wǎng)絡(luò )設計的經(jīng)驗積累,在基線(xiàn)模型的基礎上,提出了一個(gè) Efficient HDR 網(wǎng)絡(luò ),包括高效的多幀對齊和特征提取模塊兩個(gè)模塊,同時(shí)優(yōu)化了模型的訓練方法。
?。?)在多幀對齊模塊,采用 Pixel Unshuffle 操作在增大感受野的同時(shí)減少了特征圖的大小,大幅減少了后續的計算量。同時(shí),采用深度可分離卷積替代對齊模塊中的普通卷積,大幅提高運算效率。
?。?)在特征提取模塊,采用深度可分離卷積替代普通卷積,SiLU 激活函數替代 ReLU,設計了一個(gè)高效殘差特征蒸餾模塊(Efficient RFDB)。另外,探索了網(wǎng)絡(luò )深度與通道數目之間的關(guān)系,在限定計算量下層數更深且通道數少的特征提取網(wǎng)絡(luò ),可以獲得更高的評價(jià)指標。
?。?)在訓練方法上,在常規的 128x128 圖片輸入 L1 Loss 訓練后,采用了 256x256 更大尺寸輸入 + L2 Loss 進(jìn)行訓練調優(yōu)。最后,使用基于 SwinIR 搭建的 Transformer 模型作為 Teacher 模型,對前述 CNN 模型進(jìn)行蒸餾,結合 CNN 和 Transformer 各自的優(yōu)勢進(jìn)一步提升模型效果。
網(wǎng)絡(luò )結構
網(wǎng)絡(luò )的整體結構基于官方提供的 baseline 模型 AHDRNet 進(jìn)一步大幅改進(jìn)和優(yōu)化,主要可以分成三個(gè)部分:多幀對齊模塊、特征提取模塊和圖像重建模塊?;诒敬伪荣惖挠嬎懔磕繕丝紤],對網(wǎng)絡(luò )部分做了以下設計:
1. Pixel Shuffle 層:在多幀對齊模塊中使用 Pixel Unshuffle 操作(Pixel Shuffle 的逆操作),在不增加計算量的同時(shí)增大了感受野。在圖像重建模塊中使用 Pixel Shuffle 替代 AHDRNet 中的卷積操作,節省計算量。
2. 深度可分離卷積:在多幀對齊模塊和特征提取模塊,采用 Depthwise+1x1 卷積的組合替換了網(wǎng)絡(luò )中的絕大多數卷積。
3. 特征提取基礎模塊替換:在特征提取模塊采用 RFDB+ESA 替換 AHDRNet 中的 DRDB,并采用 SiLU 激活函數替換 ReLU。
4. 深度 vs 寬度:在限定計算量下,平衡特征提取模塊深度和寬度取得更好的效果。以 Track1 的約束為標準,所有模型的計算量都在 190G~200G 之間(即更深的網(wǎng)絡(luò )意味著(zhù)更少的通道數)。
整體的網(wǎng)絡(luò )結構圖如圖 2 所示,Efficient RFDB 的結構圖如圖 3 所示:
網(wǎng)絡(luò )結構圖
Efficient RFDB 結構圖
最后提交的 Track1 和 Track2 模型均采用了上述的模型結構,區別是使用 Efficient RFDB 層數和通道數有所不同,對 Track1,Efficient RFDB 層數和通道數目較多,計算量是 198.47GMACs。對 Track2,Efficient RFDB 層數相近,通道數更少,計算量是 74.02GMACs。
訓練過(guò)程
目前 HDR 的論文或比賽的主要評價(jià)指標是 PSNR-μ,指先對輸出圖片和標簽圖片分別做色調映射 tonemapping 操作再計算其 PSNR:
其中,
是指 tonemapping 操作。對圖像做了以下處理:
其中,
。
主流的方法是對色調映射后的圖片求 L1 Loss,譬如去年 NTIRE 比賽的多幀 HDR 比賽冠軍 ADNet 是使用 tonemapped 后的圖片損失函數,即:
該研究訓練過(guò)程第一步同上使用 L1 Loss,并在此基礎上加了后續三個(gè) finetune 的過(guò)程,按順序分別是:
1. L2 Loss finetune:為了獲得更高的評價(jià)指標,在微調階段該研究采用了和 PSNR 計算一致的 L2 Loss 代替 L1 Loss:
2. 大尺寸圖片 + L2 Loss finetune:由于最后用了深層的網(wǎng)絡(luò )設計,網(wǎng)絡(luò )具有更大的感受野,采用 256x256 替代 128x128 的大尺寸圖片進(jìn)行微調可以讓模型取得更好的效果。
3. 知識蒸餾 + 大尺寸 + L2 Loss finetune:使用基于 Transformer 的 SwinIR 中的 RSTB 代替 Efficient RFDB 搭建 Teacher 網(wǎng)絡(luò ),此階段損失函數如下:
其中,TS 表示 Teacher Surpervision,DS 表示 Data Supervision,最后在實(shí)驗中
實(shí)驗結果
賽道 1 中,網(wǎng)易互娛 AI Lab(ALONG)提出的方法在 PSNR-μ和 PSNR 上均是第一。如表 1 所示,主要評價(jià)指標 PSNR-μ比第二名高出了 0.172,而第二到第四的 PSNR-μ差距僅為 0.089,相比第五名之后的隊伍更是拉開(kāi)了 0.45 以上的差距。
賽道 2 中,網(wǎng)易互娛 AI Lab(ALONG)提出的方法取得了最低的計算量(GMACs)和最少的參數量(Param)。如表 2 所示,在超過(guò)基線(xiàn)方法 PSNR 和 PSNR-μ的基礎上,減少了約 40 倍的計算量。相比第二名和第三名有較大領(lǐng)先,僅使用了約一半的計算量。
表 1:賽道 1(保真度賽道)結果排名
表 2:賽道 2(低復雜度賽道)結果排名