2024IJCAI | MetalISP: 仅用1M参数的RAW到RGB高效映射模型

文章标题是：《MetaISP:Effcient RAW-to-sRGB Mappings with Merely 1M Parameters》

MetaISP收录于2024IJCAI，是新加坡国立大学（Xinchao Wang为通讯作者）和华为联合研发的新型ai-isp。

原文链接：MetaISP

【1】论文的主要motivation为目前的ai-isp为了提高在不同场景下的泛化性，一味的提增加AI-ISP架构的尺寸和复杂性，导致得到的模型根本无法在侧端（手机端）部署。

【2】论文的主要创新为：提出了一个可以针对不同输入自适应修改"parameters and network architecture"的新型ai-isp。这一点主要是依靠两个模块来实现的。MCCB（Meta Channel Correction Block）学习特征图中层级关系、MSCB（Meta Spatial Correction Block）学习空间注意力图，从而自适应的调整特征图包含的内容。

1. Abstract

SOTA的深度ISP模型通过增加网络的尺寸和复杂性缓解了对不同异质输入的有限泛化能力问题，但也不可避免地导致了网络参数量和FLOPs的急剧增长。为了解决这个问题，本文提出了MetaISP，一个通过"针对不同输入自适应调整参数和结构"以达到优越的图像重构质量的模型。我们的理论主要涉及为不同的输入生成对应的空间/通道注意力矫正矩阵，以帮助分配注意力（即模型应更关注图像的哪些部分）。我们通过“为每张输入图像预测动态权重并把权重与多个可学习的基础矩阵相结合以构建矫正矩阵“的方法实现上述理论。MetaISP在达到最好效果的同时保证了计算高效性。MetaISP在两个大规模数据集上达到了SOTA，在ZRR数据集上，MetaISP达到了23.8dB的PSNR（以9.2%的参数量和10.6%的FLOPs超过了之前的SOTA方法0.19dB）；在MAI21数据集上，MetaISP达到了25.06dB（以0.9%的参数量和2.7%的FLOPs超过了之前的SOTA方法0.17dB）

2. Methods

先验知识：self-attention的计算复杂度很高，大概是 O(N2)，N是输入序列长度；以SE (Squeeze and Excitation) 模块为代表的一系列基于CNN实现通道注意力机制的方法计算复杂度就比较低，大约是 O(r2) ，r是放缩率。

模型框架图如图1所示

图1. MetalISP的整体框架

 

2.1 整体流程

采用U-Net架构保证了模型可以学习到图像特征的底层及高层特征，促进了鲁棒的表示能力。输入是4通道的RAW图 $I_{raw}\in {\mathbb{R}}^{4\times H/2\times W/2}$ ，输出是3通道的标准RGB图像 $I_{rgb}\in {\mathbb{R}}^{3\times H\times W}$ 。MFEB将$I_{raw}$提取成 $F_0\in {\mathbb{R}}^{C\times H/2\times W/2}$ 。

到U-Net架构的bottom层之前是4对MCCB和MSCB的堆叠，生成的特征图通道数为 { C , C , C , 2 C } \{C,C,C,2C\} {C,C,C,2C}，低通道数意味着卷积核的数量也很少，保证了低计算复杂度。图中的下采样和上采样均使用的是3×3的卷积核和离散小波变换 。残差连接是基于逐像素加法实现的。

2.2 Meta Channel Correction Block (MCCB)

作者在介绍MCCB之前，先提了一嘴2018年很火的通道注意力 。作者说，这篇工作被广泛用于各个low-level的任务中 (包括深度ISP)去捕捉通道间的依赖关系，从而关注富有信息的特征并抑制冗余的特征。通道注意力CA(Channel attention)，实现了计算高效性并且把全局信息带入了特征图中。但是，作者认为，CA中使用的计算方式（通道乘法及注意力预测层）过于简单，因此会导致次优的结果。与此同时，作者提到，基于Self-attention改进的算法又面临计算复杂度$O(N^2)$ 过大的问题，因此，作者提出了MCCB模块。

考虑输入的特征图经过层归一化，卷积，GELU之后，得到特征图$X\in {\mathbb{R}}^{C\times H\times W}$被分别投射到高维空间$Y\in {\mathbb{R}}^{K\times H\times W}$（获得更丰富的特征）以及预测得到权重向量 { w n } n = 1 , 2 , . . . , N = p ( X ) \{w_n\}_{n=1,2,...,N}=p(X) {wn​}n=1,2,...,N​=p(X)。矫正矩阵（私以为可以理解为注意力图）的计算方式则为 $M={\Sigma }_{n=1}^N{w}_n(x)\varphi (n)$ ，其中， { ϕ n } n = 1 , 2 , . . . , N \{ϕ_n\}_{n=1,2,...,N} {ϕn​}n=1,2,...,N​为基础矩阵。

后面的流程就不需赘述了，一切如流程图所示。

需要注意的是，MCCB仅仅使用了一些基础矩阵就修正了特征图并应用了动态软权重实现内容自适应的通道变换。MCCB在利用了CA的低计算复杂度的同时，实现了对通道间复杂关系的高效学习。

2.3 Meta Spatial Correction Block (MSCB)

作者在讲MSCB之前，先提了一嘴基于CNN的空间注意力机制。作者指出，使用CNN去学习注意力图强调了相关的空间区域，并且利用参数共享机制降低了计算复杂度，但是，也造成了有限的感受野，阻碍了其建模长范围的像素交互。然后勾连了上一节的SA叙述 (SA可以很好的感受全局信息但是计算复杂度很高，是输入序列长度的平方次），指出SA不适用于ISP任务中的高分辨率图像。

为了利用CNN的低计算复杂度以及SA的全局上下文建模能力，作者提出了MSCB。

如图1所示，MSCB和MCCB的基础架构其实差不多，矫正向量的公式为： $V={\Sigma }_{n=1}^N{w}_n(x){\epsilon }_n$ ，其中 ${\epsilon }_n$ 是第n个basis vector。

接下来有一些很细节的点，我们放大一下MSCB看看。

图2. MSCB

 

令改变形状后的特征图为$Y$，则Y中的每一行都对应一个通道的特征图，$Y$中的每一竖列都对应一个像素在不同通道中的取值。接下来，作者使用公式$S_{w,h}=\frac{Yw,h\ast V}{|Y_{w,h}||V|}$来得到空间注意力图，再用逐元素乘法得到输出的$\hat{Y}=Y\otimes S$ ，此处，我理解就是对每个元素使用矫正向量进行矫正。作者指出，这个过程的计算复杂度为$O(HW)$，是线性复杂度，同时高效利用了特征图的全局信息。

图3. MSCB的注意力可视化

 

作者可视化了MSCB的注意力图，并指出，U-Net架构不同层级的注意力关注点不同。

3. Experiments

贴一下ZRR数据集的可视化结果图，在颜色上确实贴近GroundTruth.

图4. 实验结果可视化  

【思考1】刚看完摘要的时候，我在想：针对不同输入自适应调整参数和结构是啥意思？难道网络权重在测试的时候还可以改变？看到Methods才明白，论文是搞了一些基础矩阵（Basis Matrices），不同的图像产生不同的权重，和矩阵运算得到注意力图，再跟对应层的特征运算，得到修正后的特征（这里以MCCB为例，MSCB也差不多）。这样相当于是对特征做了处理，但是这个处理模块对于不同的图像是不同的，此之谓"动态参数和动态结构"。

【思考2】作者在MCCB和MSCB架构中是怎么同时利用了CNN和SA的？依照我个人的理解，作者其实相当于简化了生成注意力图的过程，他们设计了一些basis matrices/vectors，只需用CNN给出weights就能生成注意力图，这个过程在SA中是$O(N^2)$ 复杂度的。