[论文解读] ISTA-Net++: Flexible Deep Unfolding Network for Compressive Sensing

Di You, Jingfen Xie+1 more|arXiv (Cornell University)|Mar 22, 2021

Sparse and Compressive Sensing Techniques参考文献 26被引用 8

一句话总结

ISTA-Net++ 提出了一种用于压缩感知的灵活深度展开网络，通过条件模块动态适应多种采样率，同时采用跨块策略和平衡的训练数据集（Train400），有效减少块效应并提升在多样化图像场景下的鲁棒性。该方法在四个数据集上均达到最先进性能，在多采样率与多场景任务中保持一致的性能增益。

ABSTRACT

While deep neural networks have achieved impressive success in image compressive sensing (CS), most of them lack flexibility when dealing with multi-ratio tasks and multi-scene images in practical applications. To tackle these challenges, we propose a novel end-to-end flexible ISTA-unfolding deep network, dubbed ISTA-Net++, with superior performance and strong flexibility. Specifically, by developing a dynamic unfolding strategy, our model enjoys the adaptability of handling CS problems with different ratios, i.e., multi-ratio tasks, through a single model. A cross-block strategy is further utilized to reduce blocking artifacts and enhance the CS recovery quality. Furthermore, we adopt a balanced dataset for training, which brings more robustness when reconstructing images of multiple scenes. Extensive experiments on four datasets show that ISTA-Net++ achieves state-of-the-art results in terms of both quantitative metrics and visual quality. Considering its flexibility, effectiveness and practicability, our model is expected to serve as a suitable baseline in future CS research. The source code is available on https://github.com/jianzhangcs/ISTA-Netpp.

研究动机与目标

• 解决现有基于深度学习的压缩感知方法在实际应用中处理多种采样率时缺乏灵活性的问题。
• 通过缓解训练数据中的数据集不平衡问题，提升模型在多样化图像场景（尤其是‘少见’场景）下的重建鲁棒性。
• 通过创新的跨块策略，减少重建图像中的块效应。
• 开发一个无需微调即可实现多采样率压缩感知重建的统一模型，提升实用性与效率。

提出的方法

• 提出一种动态展开策略（DUS），通过条件模块将采样率信息注入网络的每一阶段，实现单模型对多种压缩感知采样率的自适应。
• 采用受采样率条件控制的动态梯度下降模块（DGDM）和动态邻近映射模块（DPMM），实现在重建过程中的自适应优化。
• 应用跨块策略（CBS），实现相邻图像块之间的特征交换，以抑制块效应并提升纹理保真度。
• 在平衡数据集（Train400）上进行训练，该数据集包含除标准数据集（如Train91）中主导的‘植物’类别外的多样化场景。
• 采用三阶段框架：采样（建模为卷积层）、初始化（解决维度不匹配问题）和重建（通过结合DUS与CBS的深度展开实现）。
• 采用端到端训练，损失函数最小化重建误差，实现所有组件的联合优化。

实验结果

研究问题

• RQ1 单个深度学习模型是否能在无需微调的情况下，有效处理多种采样率的压缩感知重建？
• RQ2 数据集不平衡在多大程度上影响压缩感知模型在未见图像场景下的泛化能力？该问题是否可被缓解？
• RQ3 跨块策略在基于深度展开的压缩感知重建中，能在多大程度上减少块效应？
• RQ4 使用更平衡的训练数据集（如Train400）是否能显著提升模型在多样化场景下的重建鲁棒性？

主要发现

• 在Set11数据集上，ISTA-Net++在50%采样率下平均PSNR达到38.73 dB，相比未使用条件模块的基线模型，在10%、30%和50%采样率下平均提升0.26 dB。
• 与无CBS的基线相比，跨块策略（CBS）可将PSNR提升最高达0.64 dB，并有效减少块效应。
• 在相同训练轮次下，使用平衡数据集Train400训练的模型在30%采样率下达到34.81 dB的PSNR，优于在不平衡数据集Train91上训练的模型（34.25 dB）。
• 在DIV2K验证集30%采样率下，使用Train400训练的ISTA-Net++模型在‘房屋’场景中相比使用Train91训练的模型提升0.97 dB，表明其对罕见场景的鲁棒性显著增强。
• 动态展开与跨块策略的结合在所有测试采样率和数据集中均带来一致的性能增益，证实了两个组件的有效性。
• 实验表明，数据集平衡性对性能的影响大于数据集规模，即使在Train400基础上增加更多数据（如加入Train91），性能提升也微乎其微。