文章目录
- 一、论文信息
- 二、创新点
- 三、方法
- MAB模块解读
- 1、MAB模块概述
- 2、MAB模块组成
- 3、MAB模块的优势
- 四、效果
- 五、实验结果
- 六、总结
- 代码
一、论文信息
- 标题: Multi-scale Attention Network for Single Image Super-Resolution
- 作者: Yan Wang, Yusen Li, Gang Wang, Xiaoguang Liu
- 机构: 南开大学
- 发表会议: CVPR 2024 Workshops
- 论文链接: arXiv
- GitHub代码库: icandle/MAN
二、创新点
-
多尺度大核注意力（MLKA）: 结合了多尺度机制与大核卷积，能够有效捕捉不同尺度的信息，避免了传统方法中常见的“块状伪影”问题。
-
门控空间注意力单元（GSAU）: 通过引入门控机制，优化了空间注意力的计算，去除了不必要的线性层，从而提高了信息聚合的效率和准确性。
-
灵活的网络结构: 通过堆叠不同数量的MLKA和GSAU模块，构建出多种复杂度的网络，以实现性能与计算量之间的平衡。
三、方法
-
网络架构: Multi-scale Attention Network (MAN)由三个主要模块组成：
- 浅层特征提取模块(SF): 负责初步的特征提取。
- 深层特征提取模块(DF): 基于多个多尺度注意力块（MAB），进一步提取丰富的特征。
- 高质量图像重建模块: 将提取的特征用于最终的图像重建。
-
多尺度注意力块(MAB):
- MLKA模块: 结合大核注意力、多个尺度机制和门控聚合，建立不同尺度之间的相关性。
- GSAU模块: 整合空间注意力和门控机制，简化前馈网络。
-
MLKA的功能:
- 大核注意力: 通过分解卷积建立长距离关系。
- 多尺度机制: 增强固定LKA以学习全尺度信息的注意力图。
- 门控聚合: 动态调整注意力图以避免伪影。
MAB模块解读
1、MAB模块概述
MAB（Multi-scale Attention Block）是Multi-scale Attention Network (MAN)中的核心组件，旨在通过结合多尺度大核注意力（MLKA）和门控空间注意力单元（GSAU）来提升图像超分辨率的性能。MAB模块的设计旨在有效捕捉图像中的局部和全局特征，同时避免传统卷积网络中常见的“块状伪影”问题。
2、MAB模块组成
MAB模块主要由以下两个部分构成：
-
多尺度大核注意力（MLKA）:
- 功能: MLKA通过引入多尺度机制，结合大核卷积，能够在不同尺度上提取丰富的特征信息。
- 结构:
- 首先，MLKA使用点卷积（Point-wise convolution）调整通道数。
- 然后，将特征分成三组，每组使用不同尺寸的大核卷积（如7×7、21×21、35×35）进行处理，膨胀率分别设置为（2，3，4）。
- 为了避免膨胀卷积带来的“块状伪影”，MLKA引入了门控聚合机制，通过逐元素乘法将深度卷积的输出与对应组的LKA输出相结合，从而动态调整注意力图的输出。
-
门控空间注意力单元（GSAU）:
- 功能: GSAU旨在增强特征表示能力，通过结合空间注意力和门控机制，优化信息聚合过程。
- 结构:
- GSAU通常由两个分支组成，其中一个分支使用深度卷积对特征进行加权，另一个分支则通过空间自注意力机制捕捉空间上下文信息。
- 这种设计减少了不必要的线性层，降低了计算复杂度，同时增强了特征的表达能力。
3、MAB模块的优势
-
多尺度特征提取: 通过MLKA，MAB能够在多个尺度上提取特征，增强了模型对不同图像细节的敏感性。
-
减少伪影: 通过门控聚合机制，MAB有效地减少了由于膨胀卷积引起的块状伪影，提升了图像重建的质量。
-
高效的特征表示: GSAU的引入使得模型能够更好地聚合空间信息，提升了特征的表达能力，进而提高了超分辨率的效果。
四、效果
-
性能提升: 实验结果表明，MAN在多个超分辨率基准测试中表现优异，能够与当前最先进的模型（如SwinIR）相媲美，同时在计算效率上也有显著改善。
-
避免伪影: 通过MLKA和GSAU的结合，模型有效减少了图像重建中的伪影现象，提升了视觉效果。
五、实验结果
- 基准测试: 论文中使用了多个数据集（如Set5、Set14、BSD100、Urban100、Manga109）进行测试，结果显示MAN在PSNR和SSIM指标上均优于传统的超分辨率模型，尤其是在高倍数放大（如×4）时表现突出。
| 数据集 | 上采样因子 | MAN PSNR (dB) | 与SwinIR比较 |
|---|---|---|---|
| Set5 | ×2 | 38.42 | 相当 |
| ×3 | 34.91 | 略低 | |
| ×4 | 32.87 | 良好 | |
| Set14 | ×2 | 34.44 | 相近 |
| ×3 | 30.92 | 略低 | |
| ×4 | 29.09 | 良好 | |
| BSD100 | ×2 | 32.53 | 相近 |
| ×3 | 29.65 | 略低 | |
| ×4 | 27.90 | 良好 | |
| Urban100 | ×2 | 33.80 | 相近 |
| ×3 | 34.45 | 略低 | |
| ×4 | 33.73 | 良好 | |
| Manga109 | ×2 | 40.02 | 相近 |
| ×3 | 35.21 | 略低 | |
| ×4 | 31.22 | 良好 |
六、总结
Multi-scale Attention Network (MAN)通过结合多尺度大核注意力和门控空间注意力机制，成功提升了单幅图像超分辨率重建的性能和效率。该研究不仅解决了传统方法中的一些局限性，还为未来的超分辨率模型设计提供了新的思路。MAN在多个基准测试中的优异表现，证明了其在实际应用中的潜力，尤其是在需要高质量图像重建的场景中。
代码
# -*- coding: utf-8 -*-
import torch
import torch.nn as nn
import torch.nn.functional as Fclass LayerNorm(nn.Module):r""" LayerNorm that supports two data formats: channels_last (default) or channels_first.The ordering of the dimensions in the inputs. channels_last corresponds to inputs withshape (batch_size, height, width, channels) while channels_first corresponds to inputswith shape (batch_size, channels, height, width)."""def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):super().__init__()self.weight = nn.Parameter(torch.ones(normalized_shape))self.bias = nn.Parameter(torch.zeros(normalized_shape))self.eps = epsself.data_format = data_formatif self.data_format not in ["channels_last", "channels_first"]:raise NotImplementedErrorself.normalized_shape = (normalized_shape,)def forward(self, x):if self.data_format == "channels_last":return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)elif self.data_format == "channels_first":u = x.mean(1, keepdim=True)s = (x - u).pow(2).mean(1, keepdim=True)x = (x - u) / torch.sqrt(s + self.eps)x = self.weight[:, None, None] * x + self.bias[:, None, None]return xclass SGAB(nn.Module):def __init__(self, n_feats, drop=0.0, k=2, squeeze_factor=15, attn='GLKA'):super().__init__()i_feats = n_feats * 2self.Conv1 = nn.Conv2d(n_feats, i_feats, 1, 1, 0)self.DWConv1 = nn.Conv2d(n_feats, n_feats, 7, 1, 7 // 2, groups=n_feats)self.Conv2 = nn.Conv2d(n_feats, n_feats, 1, 1, 0)self.norm = LayerNorm(n_feats, data_format='channels_first')self.scale = nn.Parameter(torch.zeros((1, n_feats, 1, 1)), requires_grad=True)def forward(self, x):shortcut = x.clone()# Ghost Expandx = self.Conv1(self.norm(x))a, x = torch.chunk(x, 2, dim=1)x = x * self.DWConv1(a)x = self.Conv2(x)return x * self.scale + shortcutclass GroupGLKA(nn.Module):def __init__(self, n_feats, k=2, squeeze_factor=15):super().__init__()i_feats = 2 * n_featsself.n_feats = n_featsself.i_feats = i_featsself.norm = LayerNorm(n_feats, data_format='channels_first')self.scale = nn.Parameter(torch.zeros((1, n_feats, 1, 1)), requires_grad=True)# Multiscale Large Kernel Attentionself.LKA7 = nn.Sequential(nn.Conv2d(n_feats // 3, n_feats // 3, 7, 1, 7 // 2, groups=n_feats // 3),nn.Conv2d(n_feats // 3, n_feats // 3, 9, stride=1, padding=(9 // 2) * 4, groups=n_feats // 3, dilation=4),nn.Conv2d(n_feats // 3, n_feats // 3, 1, 1, 0))self.LKA5 = nn.Sequential(nn.Conv2d(n_feats // 3, n_feats // 3, 5, 1, 5 // 2, groups=n_feats // 3),nn.Conv2d(n_feats // 3, n_feats // 3, 7, stride=1, padding=(7 // 2) * 3, groups=n_feats // 3, dilation=3),nn.Conv2d(n_feats // 3, n_feats // 3, 1, 1, 0))self.LKA3 = nn.Sequential(nn.Conv2d(n_feats // 3, n_feats // 3, 3, 1, 1, groups=n_feats // 3),nn.Conv2d(n_feats // 3, n_feats // 3, 5, stride=1, padding=(5 // 2) * 2, groups=n_feats // 3, dilation=2),nn.Conv2d(n_feats // 3, n_feats // 3, 1, 1, 0))self.X3 = nn.Conv2d(n_feats // 3, n_feats // 3, 3, 1, 1, groups=n_feats // 3)self.X5 = nn.Conv2d(n_feats // 3, n_feats // 3, 5, 1, 5 // 2, groups=n_feats // 3)self.X7 = nn.Conv2d(n_feats // 3, n_feats // 3, 7, 1, 7 // 2, groups=n_feats // 3)self.proj_first = nn.Sequential(nn.Conv2d(n_feats, i_feats, 1, 1, 0))self.proj_last = nn.Sequential(nn.Conv2d(n_feats, n_feats, 1, 1, 0))def forward(self, x, pre_attn=None, RAA=None):shortcut = x.clone()x = self.norm(x)x = self.proj_first(x)a, x = torch.chunk(x, 2, dim=1)a_1, a_2, a_3 = torch.chunk(a, 3, dim=1)a = torch.cat([self.LKA3(a_1) * self.X3(a_1), self.LKA5(a_2) * self.X5(a_2), self.LKA7(a_3) * self.X7(a_3)],dim=1)x = self.proj_last(x * a) * self.scale + shortcutreturn x# MABclass MAB(nn.Module):def __init__(self, n_feats):super().__init__()self.LKA = GroupGLKA(n_feats)self.LFE = SGAB(n_feats)def forward(self, x, pre_attn=None, RAA=None):# large kernel attentionx = self.LKA(x)# local feature extractionx = self.LFE(x)return xif __name__ == "__main__":dim=96 # 通道要被3整除# 如果GPU可用，将模块移动到 GPUdevice = torch.device("cuda" if torch.cuda.is_available() else "cpu")# 输入张量 (batch_size, channels,height, width)x = torch.randn(2,dim,40,40).to(device)# 初始化 MAB模块block = MAB(dim)print(block)block = block.to(device)# 前向传播output = block(x)print("输入:", x.shape)print("输出:", output.shape)
