Partial Large Kernel CNNs for Efficient Super-Resolution
论文链接:http://arxiv.org/abs/2404.11848
代码链接:https://github.com/dslisleedh/PLKSR
1、摘要
开发了一种高效的基于CNN的模型成功地整合了Transformer处理长距离依赖关系和提供实例相关权重的能力。
2、关键问题
-
尽管平均CNN的FLOP数和参数数量分别为8.1倍和9.5倍,实际上在效率上优于 Transformers,平均延迟降低了2.4倍,性能提高了 4.9 倍,降低最大 GPU内存占用率(MGO)。
-
实验发现,与多次堆叠小内核相比,使用单个大内核可以在MGO级别相似的情况下提供更快的延迟。采用单个大内核与Transformer的捕获能力更加接近与使用多个小核相比,动态卷积更能生成边缘和结构等低频特征。,虽然动态卷积提供了一种生成实例相关权重的直接方法,但作者的实验结果表明它对延迟有不利影响。
3、创新点
-
元素注意力(EA)引入到模型中,它将分配对输入特征张量中每个元素赋予单独的注意力权重,从而优化性能。
-
开发了一个基于CNN的模型,结合了Transformer 处理长距离依赖和实例相关权重的能力,优化输入特征的处理和划分,提高效率。
-
提出的模型PLKSR在四个数据集上实现了最先进的性能,并且显著减少了延迟和 MGO
-
通过视觉分析和实证测试证明PLKSR有效地捕捉到了关键超分辨率的低频特征,类似于Transformer。
4、原理
单个大内核具有较低的延迟,并且略微更高的内存使用率,如图2(b)所示。此外,在所提出的small变体模型中比较了单个大内核和堆叠
3
×
3
3 \times 3
3×3 内核的性能。如表2所示,单个大内核比堆叠的
3
×
3
3 \times 3
3×3 内核具有更低的延迟和更好的性能。
如图3所示,低分辨率图像
I
L
R
I^{LR}
ILR通过
3
×
3
3 \times 3
3×3卷积进行处理,然后依次对
N
N
N个PLK Blocks进行排序,得出与另一个
3
×
3
3 \times 3
3×3卷积生成
F
h
F_{h}
Fh。这可以用公式1来表示。
F
i
n
0
=
C
o
n
v
3
×
3
(
I
L
R
)
F
o
u
t
N
=
P
L
K
B
l
o
c
k
N
(
.
.
.
P
L
K
B
l
o
c
k
0
(
F
i
n
0
)
.
.
.
)
F
h
=
C
o
n
v
3
×
3
(
F
o
u
t
N
)
(
1
)
F^{0}_{in} = Conv_{3 \times 3}(I^{LR}) \\ F^{N}_{out} = PLKBlock^{N}(...PLKBlock^{0}(F^{0}_{in})...) \\ F_{h} = Conv_{3 \times 3}(F^{N}_{out}) \ (1)
Fin0=Conv3×3(ILR)FoutN=PLKBlockN(...PLKBlock0(Fin0)...)Fh=Conv3×3(FoutN) (1)
在另一个平行分支中,
I
L
R
I^{LR}
ILR的RGB重复了
r
2
r^{2}
r2次产生如图所示的
F
l
F_{l}
Fl在公式2中,r表示上采样因子。在PixelShuffle作[32]之后,
F
l
F_{l}
Fl成为最接近的插值图像。
F
l
=
R
e
p
e
a
t
r
2
(
I
L
R
)
,
F
l
∈
R
3
⋅
r
2
×
h
×
w
(
2
)
F_{l} = Repeat_{r^{2}}(I^{LR}), \ F_{l} \in R^{3 \cdot r^{2} \times h \times w} \ (2)
Fl=Repeatr2(ILR), Fl∈R3⋅r2×h×w (2)
最终,通过将
F
h
F_{h}
Fh与
F
l
F_{l}
Fl相加得到最终的特征图
F
F
F和来自各个路径的
F
l
F_{l}
Fl,以及通过应用PixelShuffle操作进行重建的将通道维度移动到空间维度的放大图像
I
S
R
I^{SR}
ISR。这相当于公式3,其中
h
h
h和
w
w
w表示分别是特征图的高度和宽度。
F
=
F
h
+
F
l
I
S
R
=
P
i
x
e
l
S
h
u
f
f
l
e
(
F
)
,
I
S
R
∈
R
3
×
h
⋅
r
×
w
⋅
r
(
3
)
F = F_{h} + F_{l} \\ I^{SR} = PixelShuffle(F), \ I^{SR} \in R^{3 \times h \cdot r \times w \cdot r} \ (3)
F=Fh+FlISR=PixelShuffle(F), ISR∈R3×h⋅r×w⋅r (3)
4.1. Partial Large Kernel Block
PLK 块由三个主要模块组成:双卷积通道混合器 (Doubled Convolutional Channel Mixer, DCCM) 用于局部特征提取,Partial Large Kernel 卷积 (PLKC) 用于处理长距离依赖关系,以及元素注意(EA) 模块用于实例相关调制。传入的特征
F
i
n
F_{in}
Fin到PLK块通过三个主要步骤依次处理模块,然后是最后的
1
×
1
1 \times 1
1×1卷积,随后将其添加回
F
i
n
F_{in}
Fin,如公式4所示。
F
l
o
c
a
l
=
D
D
C
M
(
F
i
n
)
,
F
c
a
t
=
P
L
K
C
(
F
l
o
c
a
l
)
,
F
a
t
t
=
E
A
(
F
c
a
t
)
,
F
o
u
t
=
F
i
n
+
C
o
n
v
1
×
1
(
F
a
t
t
)
.
(
4
)
F_{local} = DDCM(F_{in}) ,\\ F_{cat} = PLKC(F_{local}) ,\\ F_{att} = EA(F_{cat}), \\ F_{out} = F_{in} + Conv_{1 \times 1}(F_{att}). \ (4)
Flocal=DDCM(Fin),Fcat=PLKC(Flocal),Fatt=EA(Fcat),Fout=Fin+Conv1×1(Fatt). (4)
4.1.1 Double Convolutional Channel Mixer.
在Transformer中,前馈神经网络(FFN)通常用于处理通道信息。使用
X
X
X和
Y
Y
Y分别表示输入和输出,FFN可以用公式5表示。
M
L
P
p
r
o
j
MLP^{proj}
MLPproj展开通道的维度(
R
c
×
h
×
w
→
R
2
⋅
c
×
h
×
w
R^{c \times h \times w} \rightarrow R^{2 \cdot c \times h \times w}
Rc×h×w→R2⋅c×h×w),并且
M
L
P
a
g
g
MLP^{agg}
MLPagg通过减少通道恢复到其原始维度(
R
2
⋅
c
×
h
×
w
→
R
c
×
h
×
w
R^{2 \cdot c \times h \times w} \rightarrow R^{c \times h \times w}
R2⋅c×h×w→Rc×h×w)其中
c
c
c表示特征图的通道大小。
Y
=
M
L
P
a
g
g
(
G
E
L
U
(
M
L
P
p
r
o
j
(
X
)
)
)
(
5
)
Y = MLP^{agg}(GELU(MLP^{proj}(X))) \ (5)
Y=MLPagg(GELU(MLPproj(X))) (5)
最近,引入了几种变体[7,16,17,19],将卷积层合并到FFN中以增强高频细节。忽略残差跳过或门控机制,这些方法可以从根本上分为两种类型。
Y
=
M
L
P
a
g
g
(
G
E
L
U
(
C
o
n
v
3
×
3
p
r
o
j
(
X
)
)
)
(
6
)
Y = MLP^{agg}(GELU(Conv^{proj}_{3 \times 3}(X))) \ (6)
Y=MLPagg(GELU(Conv3×3proj(X))) (6)
Y
=
C
o
n
v
3
×
3
p
r
o
j
(
G
E
L
U
(
M
L
P
p
r
o
j
(
X
)
)
)
(
7
)
Y = Conv^{proj}_{3 \times 3}(GELU(MLP^{proj}(X))) \ (7)
Y=Conv3×3proj(GELU(MLPproj(X))) (7)
公式6对应于卷积通道混合器(CCM)[7],而公式7引入了倒置卷积通道混合器(ICCM)。此外,本文引入了双卷积通道混合器(DCCM),用
3
×
3
3 \times 3
3×3卷积替代所有MLP,与利用
3
×
3
3 \times 3
3×3卷积仅进行一次的现有方法一致。
Y
=
C
o
n
v
3
×
3
p
r
o
j
(
G
E
L
U
(
C
o
n
v
3
×
3
p
r
o
j
(
X
)
)
)
(
8
)
Y = Conv^{proj}_{3 \times 3}(GELU(Conv^{proj}_{3 \times 3}(X))) \ (8)
Y=Conv3×3proj(GELU(Conv3×3proj(X))) (8)
用 3 × 3 3 \times 3 3×3卷积替换所有MLP,以提高提取局部特征并进行实验确认的能力其效果比其他通道混合方法更好。在PLK块中,输入 X X X对应于 F i n F_{in} Fin,并且输出 Y Y Y对应 F l o c a l F_{local} Flocal。
4.1.2 Partial Large Kernel convolution Module.
进入部分大核卷积(PLKC)模块,输入特征图
F
l
o
c
a
l
F_{local}
Flocal根据C个通道将特征图分解为两个特征图:
F
c
o
n
v
F_{conv}
Fconv卷积和
F
i
d
F_{id}
Fid表示单位映射identity。随后,该模块在
F
c
o
n
v
F_{conv}
Fconv用一个
K
×
K
K \times K
K×K大核卷积专门生成一个大核过滤特征图,记为
F
g
l
o
b
a
l
F_{global}
Fglobal。之后,
F
g
l
o
b
a
l
F_{global}
Fglobal与
F
i
d
F_{id}
Fid通道维度相拼接,得到
F
c
a
t
F_{cat}
Fcat。其公式为
F
c
o
n
v
,
F
i
d
=
S
p
l
i
t
c
h
a
n
n
e
l
(
[
F
l
o
c
a
l
,
C
]
)
,
F
g
l
o
b
a
l
=
C
o
n
v
K
×
K
(
F
c
o
n
v
)
,
F
c
a
t
=
C
o
n
c
a
t
c
h
a
n
n
e
l
(
[
F
g
l
o
b
a
l
,
F
i
d
]
)
.
(
9
)
F_{conv}, F_{id} = Split_{channel}([F_{local}, C]), \\ F_{global} = Conv_{K \times K}(F_{conv}) , \\ F_{cat} = Concat_{channel} ([F_{global, F_{id}}]). \ (9)
Fconv,Fid=Splitchannel([Flocal,C]),Fglobal=ConvK×K(Fconv),Fcat=Concatchannel([Fglobal,Fid]). (9)
4.1.3 Element-wise Attention Module.
元素注意力机制(EA)是一种注意力机制[24] ,它保持空间和通道维度尺寸不变,将其与其他注意力机制(如空间/通道注意力)区分开来,这些机制通常涉及降维[34–36]。作者利用这一点与Transformers类似的实例相关调制模块,细节如公式10。
F
a
t
t
=
F
c
a
t
⊙
S
i
g
m
o
i
d
(
C
o
n
v
3
×
3
(
F
c
a
t
)
)
.
(
10
)
F_{att} = F_{cat} \odot Sigmoid(Conv_{3 \times 3} (F_{cat})). \ (10)
Fatt=Fcat⊙Sigmoid(Conv3×3(Fcat)). (10)
从经验上发现利用EA只能提高更深层模型的性能,因此在tiny变体中将其排除在外。
5、实验
5.1. Implementation Details
PLKSR模型由28个PLK块和64个通道,采用PLKC配置16个通道(C),内核大小(K)为17。PLKSR‑tiny模型缩小到12个PLK块和64个通道,PLKC为16C和 13K,不带EA模块利用更多区块。
5.2 Training Details
PLKSR是在DIV2K据集上行的 [2] ,而PLKSR-tiny有版本,分别在DIV2K和DF2K(DIV2K + Flicker2K [37])据集上进行。使用L1损失函数训练模型,使用 96 × 96 96 \times 96 96×96大小的64个pathes,使用随机水平翻转pathes和旋转作为数据增强方法。使用Adam [38]优化器,其中 β 1 = 0.9 \beta_{1} = 0.9 β1=0.9, β 2 = 0.99 \beta_{2} = 0.99 β2=0.99,模型进行450k次迭代。初始学习率设置为2e‑4,随后在[100k, 200k,300k,400k,425k]次迭代时依次减半。 P L K S R × 3 PLKSR \times 3 PLKSR×3 和 P L K S R × 4 PLKSR \times 4 PLKSR×4根据先前的研究[2, 12, 20],使用预训练的 P L K S R × 2 PLKSR \times 2 PLKSR×2进行微调,而PLKSR‑tiny则从头开始训练。为了进行微调,我们使用 β 1 = 0.9 \beta_{1} = 0.9 β1=0.9和 β 2 = 0.99 \beta_{2} = 0.99 β2=0.99的Adam优化器来训练模型50k次迭代。微调的初始学习率设置为为2e‑4。所有训练均采用FP16精度进行加速。
5.3. Quantitative Results
为了评估所提SR模型的效率,使用延迟和MGO作为主要指标。MGO的测量方法是使用PyTorch提供的torch.cuda.max_memory_allocated内存分配函数。为了进行性能评估,利用峰值信噪比 (PSNR) 和结构相似性指数度量(SSIM),在Y通道上计算YCbCr空间中,将图像边界裁剪为与缩放因子相等的因子。实验在五个广泛使用的数据集评估性能:Set5[39]、Set14[40]、BSD100[41]、 Urban100[42]和Manga109[43]。
在表3中详述的比较分析中,PLKSR使用SOTA SR方法进行评估,包括EDSR[2]、SwinIR‑light[12]、ELAN‑light[13]、SwinIR‑NG[33]、OmniSR[15]、SRFormer‑light[16]、DLGSANet‑light[19]和DITN[20]。在缩放因子为 × 2 \times 2 ×2时,PLKSR同时在Set5、Set14 和 BSD100数据集表现出与ELAN‑light相似的延迟。值得注意的是,SRFormer‑light在 B S D 100 × 2 BSD100 \times 2 BSD100×2上,PLKSR实现了相同的性能,显著降低延迟,MGO(分别高达68.5%和86.1%)。在规模PLKSR的性能优于ELAN‑light,延迟减少43.3%,同时记录Set5、BSD100和Urban100上的最高PSNR值。虽然有些模型不使用预训练策略,PLKSR在性能和效率方面均优于SOTA模型,证实了PLKSR的卓越效率与性能的权衡。
为了突出PLKSR的效率和可扩展性,将PLKSR‑tiny与采用大核卷积[6]或MHSA(类似)机制[7, 20] 的SOTA轻量级SR模型进行了评估。如表4所示PLKSR‑tiny的表现优于其竞争对手,在Manga109数据集上以 × 2 \times 2 ×2的缩放比例下PSNR高于39,却展示了第二低的延迟和MGO。令人印象深刻的是,在缩放系数为 × 4 \times 4 ×4时,PLKSR‑tiny表现出与DITN‑real相比,延迟降低了28%,并实现了在所有评估数据集上表现最佳。这表明PLKC具有可扩展性,并且是处理长距离依赖关系的最合适实现轻量级模型。
5.4. Ablation Study
消融研究评估了内核大小的影响,通道混合器选择,以及EA的集成。如表5所示,研究揭示了内核大小的增加与性能的提高之间存在正相关性,凸显了更大卷积核的好处。在类似延迟下评估的各种通道混频器中,DCCM 脱颖而出,成为更优的选择,可提供最高性能。此外,采用EA模块实现了性能和延迟,展示了作者设计选择的有效性。
5.5. Comparison on Mobile device
为验证该模型在实际场景中的实用性和广泛适用性。作者评估延迟在移动设备(iPhone 12)上测试PLKSR‑tiny。将PLKSR‑tiny与其他采用大核卷积[6]或MHSA(类似)机制[7,20]的ESR模型在两种图像尺寸上进行比较。如表6所示,
PLKSR‑tiny实现了两种图像大小下的最低延迟。这些结果表明PLKC模块是处理现实世界边缘设备长期依赖的最高效的应用。
5.6. Large Kernel Analysis
MHSA可以有效捕捉低频特征,例如与捕捉高频纹理的卷积不同,卷积会将图像视为形状和边缘[44]。为了验证这一点,作者在
S
R
F
o
r
m
e
r
‑
l
i
g
h
t
×
2
SRFormer‑light \times 2
SRFormer‑light×2的MHSA特征图和
P
L
K
S
R
×
2
PLKSR \times 2
PLKSR×2的大/小核特征图上可视化了傅里叶变换后的相对对数振幅。图4可以看出,大核倾向于学习类似于MHSA的低频特征,而小核与MHSA和大内核相比,它能够捕捉高频特征。
图5的进一步可视化表明大内核捕捉结构特征的能力,例如表面轮廓,以及小内核捕捉纹理的能力物体中的细节,如头发或帽子。这意味着PLKC可以轻松捕获小内核难以捕获的特征,类似于MHSA。这解释了PLKSR在使用两个特征互补方面具有较好的表现。
5.7. Comparison on LAM
为了证明PLKSR能够实际利用PLKC捕获的长距离依赖关系来重建图像,作者引入了LAM[45],一种利用像素来重建特定区域表示范围的工具模型,并将其与SRFormer‑light[16] 进行比较。模型使用
17
×
17
17 \times 17
17×17的大内核,而 SRFormer‑light使用
16
×
16
16 \times 16
16×16窗口MHSA,它们具有类似的感受野。如图6所示,PLKSR可以使用更宽的像素范围来重建红色边界框。令人惊讶的是,PLKSR利用了与补丁相邻区域中的几乎所有像素,捕获了远离补丁区域的大部分边缘,这与之前展示PLK结构偏差的实验一致。
5.8. Visual Results
为了说明PLKSR带来了可视化质量的提升,作者将放大后的结果与其他SOTA SR模型SwinIR-light [12]、OmniSR [15]、CRAFT [18]和SRFormer-light [16]在Urban100数据集上进行视觉比较。在图7中,对红色边界框内的区域进行了详细放大,展示了每个模型在放大过程中的精度,同时显示了这些重建的PSNR。与其他方法不同的是,后者往往会产生过度平滑的纹理或无法捕捉细节,PLKSR成功重建了图像并捕捉到了边缘。这强烈支持了PLKSR比其他SR模型产生更具视觉吸引力的结果的事实。
6、总结
这项研究表明,与普遍认为的卷积神经网络相比Transformer效率低的观点相反,根据直接指标,实际上它们更加高效。通过将Transformer的优势融入卷积神经网络,利用 17 × 17 17 \times 17 17×17卷积来实现大的感受野和元素级别的注意力来实现上下文相关的权重,从而减少了计算开销。PLKSR在四个数据集上实现了 × 4 \times 4 ×4尺度的最先进性能,比ELAN-light在延迟上减少了42%。实验表明,PLKC在现代GPU和边缘设备上的表现与MHSA相似,但效率显著更高。