STF-顶会图像压缩方法

✨✨ 欢迎大家来访Srlua的博文（づ￣3￣）づ╭❤～✨✨

🌟🌟 欢迎各位亲爱的读者，感谢你们抽出宝贵的时间来阅读我的文章。

我是Srlua小谢，在这里我会分享我的知识和经验。🎥

希望在这里，我们能一起探索IT世界的奥妙，提升我们的技能。🔮

记得先点赞👍后阅读哦~ 👏👏

📘📚 所属专栏：传知代码论文复现

欢迎访问我的主页：Srlua小谢 获取更多信息和资源。✨✨🌙🌙

​​

​​

目录

背景

相关工作

提出的方法

窗口注意力机制

基于CNN的架构

基于Transformer的架构

实验

核心代码解读

代码部署和复现流程

---

 本文所有资源均可在该地址处获取。

本方法选自CVPR2022"The Devil Is in the Details: Window-based Attention for Image Compression"，是一种性能极高的图像压缩方法，源码可参考附件文件，同时本文会详细介绍复现过程

背景

随着视觉应用的日益增多，图像压缩已经成为图像处理领域的一个重要研究课题。传统的图像压缩方法，如JPEG、JPEG2000等，主要依赖于手工设计的规则，虽然这些方法在一定程度上解决了图像存储和传输的问题，但在处理复杂纹理和细节方面存在局限性。近年来，基于卷积神经网络（CNN）的学习型图像压缩方法展示了优越的率失真性能。然而，CNN在捕捉局部冗余和非重复纹理方面仍存在不足，这限制了图像重建质量的进一步提升。

本文通过结合局部注意力机制和全局特征学习，提出了一种新的图像压缩方法，名为“Symmetrical TransFormer (STF)”框架，并证明了其在压缩图像时的优越性能。

相关工作

在图像压缩领域，学习型图像压缩方法近年来发展迅速，基于变分自编码器（VAE）的模型在率失真性能方面优于传统的有损压缩方法。VAE通过使用线性和非线性参数分析变换将图像映射到潜在编码空间，并结合熵估计模块精确预测潜在变量的分布，从而提高了压缩效率。自回归先验和高斯混合模型（GMM）进一步增强了熵估计模块的性能，但同时也增加了计算复杂性和时间开销。

注意力机制通过模拟生物观察的内部过程，分配更多的注意力资源到关键区域，从而获得更多细节并抑制其他无用信息。在图像压缩中，非局部注意力机制已经被证明可以通过生成隐式重要性掩码来引导潜在特征的自适应处理。尽管如此，非局部注意力机制并未改变CNN结构的全局感知特性，而Transformer架构通过利用注意力机制捕捉全局依赖关系，在图像分类和语义分割等任务中取得了优异的性能。

提出的方法

窗口注意力机制

大多数先前的方法使用基于全局感受野的注意力机制生成注意力掩码。然而，在图像压缩任务中，全局语义信息的作用不如局部空间邻近元素的相关性大。为此，本文提出了窗口注意力机制，具体方法如下：

1. 局部窗口中的注意力计算：为了有效地建模并关注空间邻近元素，我们将特征图划分为非重叠的M×M窗口，在每个窗口内分别计算注意力掩码。
2. 窗口注意力模块：WAM将非局部注意力模块替换为窗口注意力模块，以更好地关注高对比度区域。通过可视化注意力模块的效果，可以看到WAM在复杂区域（高对比度）分配了更多比特，而在简单区域（低对比度）分配了较少比特，从而在这些区域内保留更多的细节。

图1 窗口注意力的可视化结果

基于CNN的架构

在现有的CNN架构中插入窗口注意力模块，能够更合理地分配比特。该模块作为一个即插即用的组件，能够显著提升现有模型的率失真性能，同时计算开销可以忽略不计。具体架构如下图所示：

图2 基于CNN的架构

1. 编码器与解码器：分别在编码器和解码器中插入WAM模块，帮助在不同区域内部合理分配比特。
2. 卷积层与残差块：在编码器和解码器中使用卷积层和残差块（Residual Block, RB）以提高特征提取和重建能力。

基于Transformer的架构

受到Transformer架构在计算机视觉任务中成功的启发，本文设计了一个对称的Transformer框架，具体步骤如下：

重新思考Transformer的设计：由于目标是验证自注意力层和MLP是否能在学习型图像压缩任务中达到与原始CNN架构相当的性能，本文设计了一个全新的对称Transformer框架，在编码器和解码器中均不使用卷积层。主要难点在于：

• 传统工作大多基于CNN消除空间冗余并捕捉空间结构，直接将图像划分为patch可能导致patch内部的空间冗余。
• GDN是图像压缩中最常用的归一化和非线性激活函数，但在深层Transformer架构中不稳定，且与注意力机制不兼容。
• 计算大视野上的注意力图并非最佳。

1. 编码器设计：将原始图像x划分为N大小的patch，通过线性嵌入层生成特征图fp。然后，将特征图重塑为序列fs，输入到Transformer块和patch合并层中，逐渐降低特征的分辨率并加倍特征通道数。

2. 解码器设计：设计对称的解码器，由多个Transformer块、patch