(即插即用模块-Attention部分) 六十一、(2024 ACCV) LIA 基于局部重要性的注意力

paper：PlainUSR: Chasing Faster ConvNet for Efficient Super-Resolution

Code：https://github.com/icandle/PlainUSR

---

1、Local Importance-based Attention

现有空间注意力机制的缺陷：1-order 注意力（如 ESA）： 性能较弱，无法充分利用图像信息。2-order 注意力（如 Self-Attention）： 计算复杂度高，运行速度慢，不适合轻量级 SR 模型。而这篇论文提出一种 基于局部重要性的注意力（Local Importance-based Attention），旨在保证性能的前提下，降低计算复杂度，实现高效的 2-order 信息交互。LIA d的原理主要有两点：局部重要性： 通过计算每个像素周围区域的局部重要性，识别图像中关键信息的位置。注意力图： 利用局部重要性生成注意力图，对特征图进行加权，增强重要信息，抑制无关信息。

LIA 的实现过程：

1. 局部重要性计算：使用 SoftPool 和 3x3 卷积对特征图进行下采样，扩大感受野，减少计算量。然后通过 Sigmoid 激活函数将下采样后的特征图转换为局部重要性图。
2. 注意力图生成：使用第一个通道的特征图作为门控信号，对局部重要性图进行加权。使用 Bilinear 插值将注意力图缩放到原始特征图的尺寸。
3. 特征图加权：将注意力图与原始特征图进行逐元素相乘，得到加权后的特征图。

优势：

• 性能： LIA 能够有效地捕捉图像中的关键信息，提升 SR 模型的性能。
• 效率： 相比于 2-order 注意力机制，LIA 计算复杂度更低，运行速度更快。
• 可扩展性： LIA 可以灵活地与其他网络结构结合，适用于不同的 SR 任务。

Local Importance-based Attention 结构图：

2、代码实现

import math
import torch
import torch.nn as nn
import torch.nn.functional as Fclass SoftPooling2D(torch.nn.Module):def __init__(self,kernel_size,stride=None,padding=0):super(SoftPooling2D, self).__init__()self.avgpool = torch.nn.AvgPool2d(kernel_size,stride,padding, count_include_pad=False)def forward(self, x):# return self.avgpool(x)x_exp = torch.exp(x)x_exp_pool = self.avgpool(x_exp)x = self.avgpool(x_exp*x)return x/x_exp_poolclass LIA(nn.Module):''' attention based on local importance'''def __init__(self, channels, f=16):super().__init__()f = fself.body = nn.Sequential(# sample importancenn.Conv2d(channels, f, 1),SoftPooling2D(7, stride=3),nn.Conv2d(f, f, kernel_size=3, stride=2, padding=1),nn.Conv2d(f, channels, 3, padding=1),# to heatmapnn.Sigmoid(),)self.gate = nn.Sequential(nn.Sigmoid(),)def forward(self, x):''' forward '''# interpolate the heat mapg = self.gate(x[:,:1])w = F.interpolate(self.body(x), (x.size(2), x.size(3)), mode='bilinear', align_corners=False)return x * w * gif __name__ == '__main__':x = torch.randn(4, 64, 128, 128).cuda()model = LIA(64).cuda()out = model(x)print(out.shape)