CVPR2021｜超分性能不降低，計算量降低50%，董超等人提出加速圖像超分的ClassSR

本文是深圳先進技術(shù)研究院孔祥濤、趙恒源、喬宇、董超等人在low-level領(lǐng)域關(guān)于超分網(wǎng)絡(luò)加速的一次探索。它創(chuàng)新性的將分類與超分進行了融合，根據(jù)不同子塊的復(fù)原難度自適應(yīng)選擇合適的超分分支以降低整體計算復(fù)雜度：復(fù)原難度低的平坦區(qū)域選擇復(fù)雜度低的超分分支，復(fù)原難度高的紋理區(qū)域選擇復(fù)雜度高的超分分支。在不降低超分性能的情況下，該方法可以最高可以節(jié)省50%的計算量。值得各位low-level領(lǐng)域的同學(xué)仔細研究一番。

Abstract

本文旨在加速超分模型在大圖像(2K-8K)上的推理速度。在實際應(yīng)用中，大圖像同嘗會被拆分為多個小圖進行處理?；谠撎幚磉^程，我們發(fā)現(xiàn)：不同圖像區(qū)域具有不同的復(fù)原難度，可以被不同容量的模型進行處理。直覺上來講，平滑區(qū)域要比復(fù)雜紋理區(qū)域更易于超分。

為利用該特性，我們采用合適的SR模型處理拆分后的不同子圖，即平滑區(qū)域采用簡單的小模型，紋理區(qū)域采用復(fù)雜的大模型。以此為根基，我們提出了一種新的解決方案：ClassSR，它將分類與超分組合到統(tǒng)一框架中。具體來講，它先采用ClassModule將子圖按照復(fù)原難度分到不同的類別；然后對不同類別的子圖采用不同的SRModule進行超分。這里的ClassModule就是常規(guī)分類網(wǎng)絡(luò)，而SRModule則是由不同大小的SR模型構(gòu)成。

我們更進一步引入一個帶Class-Loss與Average-Loss的分類方法以生成分類結(jié)果。通過聯(lián)合訓(xùn)練，大部分子圖將被送入更小的網(wǎng)絡(luò)。我們通過證實：在DIV8K數(shù)據(jù)集上，所提ClassSR可以幫助大部分現(xiàn)有方法(如FSRCNN、CARN、SRResNet、RCAN)節(jié)省50%FLOPs。注：這是一種通用性方法，同樣適用于其他low-level視覺任務(wù)。

上圖給出了本文所提方法處理示意圖以及其對于不同超分方法的指標-計算量的調(diào)整示意圖。本文的主要貢獻包含以下三點：

• 提出ClassSR，首個將分類與超分結(jié)合的圖像超分方案；
• 提出結(jié)合數(shù)據(jù)特性克服low-level加速問題，而且它與其他加速計算具有互補性；
• 提出一種帶兩種新穎損失的分類方案以促進子圖按照復(fù)原難度進行分類。

Method

我們首先來看一下DIV2K中不同子圖(將DIV2K拆分為 大小的子圖)的類型，見上圖與表。為評估不同子圖的復(fù)原難度，我們采用與訓(xùn)練MSResNet進行評估，按照其PSNR值進行排序。如圖所示，我們將其均勻的分為三組：simple、medium以及hard。很明顯：具有更高PSNR的子圖通常是平坦區(qū)域，而具有更低PSNR的子圖包含復(fù)雜的紋理。

我們采用不同的網(wǎng)絡(luò)處理不同類型的子圖，如上表所示。我們采用三個不同大小的FSRCNN(通過通道數(shù)控制大小)進行處理。這些模型分別在不同類型的子圖上進行訓(xùn)練。可以看到：在simple子圖上，F(xiàn)SRCNN(16)與FSRCNN(56)的性能基本沒區(qū)別；類似的，F(xiàn)SRCNN(36)可以在medium子圖上取得了與FSRCNN(56)相當(dāng)?shù)男阅?。這也就意味著：我們可以采用輕量型網(wǎng)絡(luò)處理simple子圖以節(jié)省計算量。這也就是本文所提ClassSR的初衷與核心所在。

ClassSR

上圖給出了本文所提ClassSRC的網(wǎng)絡(luò)架構(gòu)示意圖，它包含兩個關(guān)鍵性模塊：ClassModule與SRModule。ClassModule負責(zé)將輸入子圖像劃分到M個類別，而SRModule則包含M個分支(即不同復(fù)雜度的SR網(wǎng)絡(luò)) 。具體來講，大的輸入圖像LR圖像X首先被拆分為重疊子圖像 ；ClassModule對接收到的每個子圖進行分類，輸出概率向量 ；我們選擇概率最大(即J=argmax_j P_j(x_i))的那個分支SR網(wǎng)絡(luò)對當(dāng)前塊進行超分 ；最后我們組合所有超分子圖得到最終的大的SR圖像 。

ClassModule

該模塊的作用是告訴我們“輸入子圖是比較容易重建還是難以重建”。我們設(shè)計了一個簡單的分類網(wǎng)絡(luò)，包含5個卷積層、一個均值池化以及全連接層。卷積層負責(zé)提取特征，池化與全連接層負責(zé)輸出概率向量。該模塊極為輕量，故而僅帶來很少的計算量。實驗表明：如此簡單的分類網(wǎng)絡(luò)已經(jīng)足以取得令人滿意的分類結(jié)果。

SRModule

這里的SRModule是由多個不同復(fù)雜度的超分網(wǎng)絡(luò)分支構(gòu)成 。從廣義上來講，每個分支可以任意的可學(xué)習(xí)超分網(wǎng)絡(luò)。由于我們的目標在于：對現(xiàn)有超分方法(如FSRCNN、CARN等)進行加速。因此，我們采用這類輕量型網(wǎng)絡(luò)作為基線，并設(shè)置為最復(fù)雜的分支 ，其他分支則通過調(diào)整通道數(shù)來控制網(wǎng)絡(luò)復(fù)雜度。

那么每個分支的通道數(shù)該如何確定呢？一個基本原則：對應(yīng)分支應(yīng)當(dāng)取得基線網(wǎng)絡(luò)在對應(yīng)類別的性能，可參考下表中的通道與性能的對比。當(dāng)然，我們也可以通過調(diào)整網(wǎng)絡(luò)的深度來控制模型復(fù)雜度，但只要滿足上述原則即可。

Classification Method

在訓(xùn)練過程中，ClassModule將對輸入子圖按照復(fù)原難度(而非預(yù)定義標簽)進行進行分類。因此，與測試階段不同，在訓(xùn)練階段每個子圖均應(yīng)經(jīng)由M個SR分支進行處理。除此之外，為確保ClassModule可以從重建結(jié)果接收到回傳梯度，我們對每個重建子圖以及對應(yīng)的分類概率相乘得到最終的SR結(jié)果：

我們通過調(diào)整ImageLoss(即 損失)來約束重建結(jié)果y，我們可以自動得到分類概率。但是在測試階段，輸入子圖僅僅經(jīng)過最大概率的SR分支。因此，我們提出了 損失(Class-Loss)以使得最大概率接近1，進而確保y以概率1等價于超分子圖。

需要注意的是，如果我們僅僅采用ImageLoss與Class-Loss，訓(xùn)練過程很容易收斂到極值點(即最復(fù)雜分支)。為避免該現(xiàn)象，我們設(shè)計了 （Average-Loss）以約束分類結(jié)果。此即為本文所提出的一種新的分類方法。

Loss Functions

損失函數(shù)包含三個損失，即Image-Loss、Class-Loss以及Average-Loss。具體來講，Image-Loss用于確保圖像重建質(zhì)量；Class-Loss用于提升分類的有效性；Average-Loss用于確保每個SR分支可以被均等的選擇?？傮w損失函數(shù)定義如下：

在上述三個損失中，Image-Loss很容易理解，即( )；那么另外兩個損失如何定義呢？這就是接下來的要介紹的了。先來看一下Class-Loss,我們希望最大概率的類具有更高的置信度，比如 就比 更好。其定義如下:

該損失可以極大的提升不同分類結(jié)果的概率差異，進而使得最大概率值接近1。正如前面所提到的“僅僅采用Image-Loss與Class-Loss會使得子圖傾向于選擇最復(fù)雜分支”，這是因為復(fù)雜分支更容易取得更好的結(jié)果，這就是使得ClassModule失其本意，進而SRModule將退化為基線模型。為避免該現(xiàn)象發(fā)生，我們應(yīng)當(dāng)確保每個SR分支被均等的選擇。因此，我們提出了Average-Loss以約束分類結(jié)果，定義如下：

注：B表示batch size。由于統(tǒng)計數(shù)值不會進行梯度傳播，故而我們采用概率和計算子圖數(shù)量。通過該損失，訓(xùn)練階段經(jīng)過每個SR分支的子圖數(shù)量近似相同。

Training Strategy

如果我們同時從頭開始訓(xùn)練ClassModule與SRModule，性能會非常不穩(wěn)定，分類模塊很容易陷入不好的局部最小。因此，在訓(xùn)練階段，我們采用三步法：

• First：預(yù)訓(xùn)練SRModule。為預(yù)訓(xùn)練SRModule，我們采用PSNR值進行數(shù)據(jù)分類。具體來講，以預(yù)訓(xùn)練MSRResNet的PSNR進行劃分，前三分之一作為hard類，后三分之一作為simple類。然后，基于所劃分的數(shù)據(jù)集進行不同復(fù)雜度的SR分支模型的訓(xùn)練。
• Second：固定SRModule，訓(xùn)練ClassModule；
• Third：聯(lián)調(diào)ClassModule與SRModule。

Experiments

在數(shù)據(jù)集方面，我們選擇DIV2K，并按照四個尺度(0.6,0.7,0.8,0.9)下采樣生成HR圖像，這些圖像進一步下采樣四倍得到LR圖像，最后我們裁剪得到1.59M個大小為 的LR子圖，這些子圖將被均勻的劃分為三類。在訓(xùn)練階段，將通過鏡像與旋轉(zhuǎn)進一步增廣。除此之外，我們從DIV2K驗證集中挑選10張圖像(0801-0810)用于訓(xùn)練過程中的驗證。

在測試集方面，我們沒有選擇Set5、Set14這些(分辨率太小)，我從DIV8K中選擇300張圖像(1201-1500)。具體來講，前兩百分別下采樣到2K和4K，稱之為Test2K、Test4K，最后的100則稱之為Test8K。在測試階段，LR圖像裁剪為 大小，步長為28。

下表給出了FSRCNN、CARN、RCAN三種網(wǎng)絡(luò)及對應(yīng)ClassSR的性能對比。很明顯，相比基線方法，所提方法可以取得更加性能同時具有更低的計算復(fù)雜度，從70%~50%不等。計算復(fù)雜度的降低與測試數(shù)據(jù)高度相關(guān)，Test8K可以取得近50%的計算復(fù)雜度降低。這是因為更大的圖形可以拆分為更多的子圖，包含更多的simple子圖。

與此同時，下圖還提供了不同方法的視覺效果對比?？梢钥吹剑?strong style="font-weight: bold;color: black;">ClassSR與基線方案具有相同的視覺效果，且不同區(qū)域的過度平滑而自然。也就是說：不同區(qū)域的不同處理不會引入不連貫現(xiàn)象。

作者同時還提供了子圖分類結(jié)果的示意圖，如下所示。綠色塊代表simple，黃色塊代表medium，紅色塊代表hard。

全文到此結(jié)束，更多消融實驗與分析建議各位同學(xué)查看原文。原文關(guān)于ClassModule、Loss收斂、網(wǎng)絡(luò)參數(shù)配置等均有非常詳細的分析，感興趣的同學(xué)建議移步原文。