【速览】TIP 2020 |AMP-Net：一种针对图像压缩感知的深度学习型稀疏优化

原创张衷豪中国图象图形学学会CSIG

。

学会“成果速览”系列文章旨在将图像图形领域会议期刊重要成果进行传播，通过短篇文章让读者用母语快速了解相关学术动态，欢迎关注和投稿~

◆ ◆ ◆ ◆

AMP-Net：一种针对图像压缩感知的深度学习型稀疏优化

张衷豪

^1

，刘翼鹏

^1

，刘佳妮

^1

，文飞

^2

，朱策

^1

^1

电子科技大学，

^2

上海交通大学

TIP 2020

撰稿人：张衷豪

推荐理事：林宙辰

原文标题：AMP-Net: Denoising based Deep Unfolding for Compressive Image Sensing

论文地址：https://ieeexplore.ieee.org/document/9298950

论文代码：https://github.com/ZhonghaoZ/AMP-Net_TIP

◆ ◆ ◆ ◆

摘要

压缩感知在低于奈奎斯特采样率的情况下进行采样和重建。当目标信号满足一定稀疏性条件时，能够将信号无失真的重建。大多数压缩感知重建算法可以分为两类，即基于模型的方法和深度网络方法。通过将迭代优化算法展开成神经网络，深度展开模型能同时具备基于模型的方法的解释性和经典深度神经网络方法速度快的特点。

在本文中，为解决视觉图像压缩感知问题，我们提出了一种名为AMP-Net的深度展开模型。它将近似消息传递（Approximate Message Passing, AMP）算法的迭代去噪过程展开成多层网络。此外，AMP-Net集成了解块模块，以消除通常出现在视觉图像压缩感知中的块状伪影。另外，我们将采样矩阵与其它网络参数联合训练，达到采样-重建联合训练的性能。实验结果表明，AMP-Net具有比其它现有方法更高的重建精度，更快的重建速度，更少的网络参数。相关技术可以推广到其它图像降噪、重建、增强等应用。

背景

压缩感知通过采集少量信息来实现图像的快速采样，数学表达式可表示为 $\mathbf{y}=\mathbf{A}\mathbf{x}$ ，其中 $\mathbf{x}\in\mathbb{R}^N$ 是向量化的原始数据， $\mathbf{A}\in\mathbb{R}^{M \times N}$ 为采样矩阵，其中 $M< N$ ， $\mathbf{y}\in\mathbb{R}^N$ 是采样得到的样本。利用 $\mathbf{y}$ 来重建 $\mathbf{x}$ 是一个欠定的问题。一个常用的解法是采用以下优化模型： $\min_\limits{\mathbf{x}}\Re \left ( \mathbf{x} \right ),s.t.\mathbf{y}= \mathbf{A}\mathbf{x},$ 其中， $\Re \left ( \mathbf{x} \right )$ 是表示图像先验信息的正则项。

有许多算法被设计来解决这个问题。基于模型的方法选择不同的正则项，如稀疏、低秩等，并利用非线性迭代算法来求解该问题。这类方法具有可解释性，但是通常比较耗时。端到端深度神经网络（如卷积神经网络）将 $\mathbf{y}$ 直接映射至 $\mathbf{x}$ 。这一类方法重建速度快，但是不具有可靠的理论保证。

现有的深度展开模型可大致分为两类。第一类模型结合人为设计的正则项来构建深度展开模型。然而，人为设计的正则项可能会忽略一些表征数据语义的细节信息，而这一类信息对于重建是有帮助的。对此，第二类模型通过神经网络来对正则项进行学习，来获得这些信息。然而，目前第二类模型的参数量都很大，这对模型的实际应用有一定的阻碍。

本文基于近似消息传递算法的去噪形式，搭建了名为AMP-Net的深度展开网络。此外，为解决自然图像中逐块采样产生的块效应，引入了去块效应模块。并且，对采样矩阵进行了训练，来获得同时采样-重建耦合的采样方案。

AMP算法的去噪视角

AMP算法[1]分析以下两个等式 $\mathbf{z}^{k-1}=\mathbf{y}-\mathbf{A}\mathbf{x}^{k-1},$ $\mathbf{x}^{k}=\mathfrak{T}_{k}\left ( \mathbf{A}^\mathrm{T}\mathbf{z} ^{k-1}+\mathbf{x}^{k-1}\right ),$

若 $\bar{\mathbf{x}}$ 为原始数据，则以下等式成立 $\mathbf{A}^\mathrm{T}\mathbf{z}^0+\mathbf{x}^0= \bar{\mathbf{x}}+\left ( \mathbf{A}^\mathrm{T}\mathbf{A}-\mathbf{I} \right )\left (\bar{\mathbf{x}}- \mathbf{x}^0 \right )$

若采样矩阵 $\mathbf{A}$ 的元素 $\mathbf{A}_{ij}$ 服从高斯分布，即 $\mathbf{A}_{ij}\sim \mathbf{N}\left ( 0,1/M \right )$ 。则 $\left ( \mathbf{A}^\mathrm{T}\mathbf{A}-\mathbf{I} \right )\left (\bar{\mathbf{x}}- \mathbf{x}^{k-1} \right )$ 也可被证明服从高斯分布。上式即可表示为 $\mathbf{A}^\mathrm{T}\mathbf{z}^{k-1}+\mathbf{x}^{k-1}=\bar{\mathbf{x}}+\mathbf{e},$ 其中 $\mathbf{e}=\left ( \mathbf{A}^\mathrm{T}\mathbf{A}-\mathbf{I} \right )\left (\bar{\mathbf{x}}- \mathbf{x}^{k-1} \right )$ 为服从高斯分布的噪声项。则 $\mathfrak{T}_k(\cdot)$ 可被解释为去噪器，其具体形式与图像的先验信息有关。在本文中，我们将其称为AMP算法的去噪视角。

AMP-Net 的结构

AMP-Net可分为采样模型和重建模型两部分。其中重建模型由一个初始化模块和多个重建模块构成。图1展示了AMP-Net的框架。

图 1 AMP-Net的框架

1.采样模型

采样模型对图像进行逐块采样。定义函数 $\mathfrak{S}\space(\mathbf{X},n)$ ，它将图像 $\mathbf{X}\in \mathbb{R}^{L\times P}$ 分割成 $n\times n$ 大小的图像块 $\mathbf{X}_i\in \mathbb{R}^{n\times n}$ ，其中 $i\in \left \{ 1,2,\cdots ,I \right \}$ 且 $L\cdot P= I\cdot n^{2}$ 。在文章中， $n$ 被设置为33。定义函数 $\mathbf{vec}\left ( \cdot \right )$ ，它用于对图像块进行向量化，满足 $\mathbf{vec}\left ( \mathbf{X}_{i} \right )\in \mathbb{R}^{n^{2}}$ 和 $\mathbf{vec}\left ( \mathfrak{S}\left (\mathbf{X},n\right ) \right )\in \mathbb{R}^{n^{2}\times I}$ ，则采样过程可以表示为 $\mathbf{Y}=\mathbf{A}\mathbf{vec}(\mathfrak{S}(\mathbf{X},n)),$ 其中 $\mathbf{A}\in \mathbb{R}^{M\times n^{2}}$ 为采样矩阵， $\mathbf{Y}\in \mathbb{R}^{M\times I}$ 为采样得到的样本。图2展示了采样的过程。

图 2 采样模型

2.初始化模块

初始化模块将采样结果恢复至原图大小，并作为重建模型的输入。定义 $\mathfrak{S}^{-1}\left ( \cdot ,n \right )$ 和 $\mathbf{vec}^{-1}\left ( \cdot\right )$ 分别为 $\mathfrak{S}(\cdot,n)$ 和 $\mathbf{vec}^{-1}\left ( \cdot \right )$ 的逆函数。我们引入可训练的矩阵 $\mathbf{B}\in\mathbb{R}^{n^{2}\times M}$ ，并将初始化过程表示为 $\mathbf{X}^0=\mathfrak{S}^{-1}(\mathbf{vec}^{-1}(\mathbf{BY}),n)$ ，其中 $\mathbf{X}^0\in\mathbb{R}^{L\times P}$ 为初始化后的图像。图3展示了采样后的过程。

图 3 初始化模块

3.重建模块

去噪模块是基于AMP算法的去噪视角建立的。通过该视角，原始图像块通过下式获得

若能获得 $\bar {\mathbf{x}}_{i}-\mathbf{x}_{i}^{k-1}$ 的值，则能完整重建出原始图像。对此，我们引入深度神经网络来拟合 $\bar {\mathbf{x}}_{i}-\mathbf{x}_{i}^{k-1}$ ，并将上式拓展为迭代算法的形式，并引入一个控制参数 $\alpha_k$ 。第 $k$ 个重建模块的重建过程可以表示为

其中 $\mathfrak{N}\space _k(\cdot)$ 为引入的卷积神经网络。对整个图像的重建可表示为

由去噪模块的重建过程可以看出，采样矩阵在重建过程中发挥着重要的作用，若能获得一个合适的采样矩阵，则有希望提高重建效果。

此外，我们引入去块效应模块来去除逐块采样与重建产生的块效应。第 $k$ 个去块效应模块可表示为

其中 $\mathfrak B_k(\cdot)$ 被定义为卷积神经网络。图4展示了第 $k$ 个重建模块（包括一个去噪模块和一个去块效应模块）。

图 4 AMP-Net的重建模块

在本文中，AMP-Net-K是有

K

个去噪模块的AMP-Net；AMP-Net-K-B是有具有去块效应模块的AMP-Net-K；AMP-Net-K-M是采用了采样矩阵训练策略的AMP-Net-K；AMP-Net-K-BM是采用了采样矩阵训练策略的AMP-Net-K-B。文章采用重建图像和原始图像之间的均方误差（MSE）作为损失函数。

实验结果

1.验证AMP-Net的展开策略

表1和表2对比了不同方法在不同采样率情况下，在Set11和BSDS500测试集上的平均PSNR和SSIM。图5和图6展示了具有不同数量的重建模块的深度展开模型在两个不同采样率下的平均PSNR。可以看出，AMP-Net展开策略是有效的。

表 1 AMP-Net-K与其他模型在Set11上的性能

表 2 AMP-Net-K与其他模型在BSDS500测试集上的性能

图 5 采样率为30%时，不同模型在BSDS500测试集上的平均PSNR

图 6 采样率为10%时，不同模型在BSDS500测试集上的平均PSNR

2.验证去块效应模块的效果

表3展示了在Set11上，不同的模型采用不同去块效应策略的平均PSNR和SSIM。图7展示了不同的模型采用不同去块效应策略在对Monarch图像的重建效果。通过表3和图7可以看出去块效应模块有很好的效果，并且具有通用性。

表 3 不同模型使用不同的去块效应策略在Set11上的性能

图 7 采样率为10%时，不同的去块效应策略在Monarch图像上的效果

3.验证采样矩阵训练策略

表4和表5展示了不同的采样率下AMP-Net-K-M的效果，图8展示了不同模型使用不同采样矩阵的重建效果。可以看出，通过训练采样矩阵，AMP-Net的效果得到了提高，并且训练好的采样矩阵也具有通用性。

表 4 AMP-Net-K-M与AMP-Net-K在Set11上的性能

表5 AMP-Net-K-M与AMP-Net-K在BSDS500测试集上的性能

图 8 在30%采样率下，不同模型使用不同采样矩阵的性能

4.验证AMP-Net-K-BM的效果

表6和表7对比了AMP-Net-K-BM与其他模型的效果。可以看出，AMP-Net-K-BM优与其他模型。并且，在一定范围内，当重建模块数量变多时，重建效果呈现增加的趋势。

表 6 AMP-Net-K-BM与其他模型在Set11上的性能

表7 AMP-Net-K-BM与其他模型在BSDS500测试集上的性能

结论

我们基于AMP算法的去噪视角提出了一种新的用于图像压缩感知的深度展开模型AMP-Net。实验依次验证了展开策略、去块效应模块和采样矩阵训练策略的有效性。此外，实验还验证了具有去块效应模块和采样矩阵训练策略的AMP-Net的重建效果。实验结果表明，AMP-Net-K-BM相比其他方法而言，具有更好的重建性能。AMP-Net-2-BM的性能已经优与了其他对比方法，并且，当重建模块数量增加时，重建性能也随之增加。

参考文献

[1] Donoho, David L , Maleki, et al. Message-passing algorithms for compressed sensing.[J]. Proceedings of the National Academy of Sciences of the United States of America, 2009.

继续滑动看下一个