【Block总结】OutlookAttention注意力，捕捉细节和局部特征|即插即用

论文信息

标题: VOLO: Vision Outlooker for Visual Recognition
作者: Li Yuan, Qibin Hou, Zihang Jiang, Jiashi Feng, Shuicheng Yan
代码链接: https://github.com/sail-sg/volo
论文链接: https://arxiv.org/pdf/2106.13112

创新点

前景注意力机制: VOLO引入了一种称为“outlook attention”的新型注意力机制，能够动态地在输入图像上进行局部特征聚合。这种机制专注于编码细粒度特征，而不是传统自注意力机制所关注的全局依赖性，从而提高了模型在视觉识别任务中的表现。
高效的特征编码: VOLO通过滑动窗口的方式进行局部特征聚合，打破了自注意力机制在计算复杂度上的瓶颈，使得模型在内存使用上更加高效。

方法

模型架构: VOLO的架构相对简单，主要包括以下几个部分：
- Outlook Attention: 该机制通过局部窗口内的相似度计算生成注意力权重，有效地聚合细粒度特征。
- 多层感知机（MLP）: 用于进一步处理和整合特征，增强模型的表达能力。

实验结果

ImageNet-1K分类任务: VOLO在该任务中实现了87.1%的top-1准确率，成为首个在该数据集上超过87%准确率的模型，且未使用任何额外训练数据。与其他模型相比，VOLO在参数量仅为296M的情况下，表现出色，显示出其高效性。
下游任务表现: VOLO在CityScapes和ADE20K等下游任务中也表现优异，分别取得了84.3%和54.3%的mIoU（平均交并比）得分，证明了其良好的迁移学习能力。

总结

VOLO通过引入前景注意力机制和高效的特征编码方法，显著提升了视觉识别模型的性能，尤其是在细粒度特征的处理上。该模型在多个标准数据集上取得了优异的成绩，为未来的视觉识别研究提供了新的思路和方向。VOLO的设计理念和实验结果表明，基于注意力的模型在视觉识别领域具有广泛的应用潜力。

代码

import torch
import torch.nn as nn
import math
import torch.nn.functional as F

class OutlookAttention(nn.Module):
    """
    Implementation of outlook attention
    --dim: hidden dim
    --num_heads: number of heads
    --kernel_size: kernel size in each window for outlook attention
    return: token features after outlook attention
    """

    def __init__(self, dim, num_heads, kernel_size=3, padding=1, stride=1,
                 qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
        super().__init__()
        head_dim = dim // num_heads
        self.num_heads = num_heads
        self.kernel_size = kernel_size
        self.padding = padding
        self.stride = stride
        self.scale = qk_scale or head_dim**-0.5

        self.v = nn.Linear(dim, dim, bias=qkv_bias)
        self.attn = nn.Linear(dim, kernel_size**4 * num_heads)

        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

        self.unfold = nn.Unfold(kernel_size=kernel_size, padding=padding, stride=stride)
        self.pool = nn.AvgPool2d(kernel_size=stride, stride=stride, ceil_mode=True)

    def forward(self, x):
        B, H, W, C = x.shape

        v = self.v(x).permute(0, 3, 1, 2)  # B, C, H, W

        h, w = math.ceil(H / self.stride), math.ceil(W / self.stride)
        v = self.unfold(v).reshape(B, self.num_heads, C // self.num_heads,
                                   self.kernel_size * self.kernel_size,
                                   h * w).permute(0, 1, 4, 3, 2)  # B,H,N,kxk,C/H

        attn = self.pool(x.permute(0, 3, 1, 2)).permute(0, 2, 3, 1)
        attn = self.attn(attn).reshape(
            B, h * w, self.num_heads, self.kernel_size * self.kernel_size,
            self.kernel_size * self.kernel_size).permute(0, 2, 1, 3, 4)  # B,H,N,kxk,kxk
        attn = attn * self.scale
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)

        x = (attn @ v).permute(0, 1, 4, 3, 2).reshape(
            B, C * self.kernel_size * self.kernel_size, h * w)
        x = F.fold(x, output_size=(H, W), kernel_size=self.kernel_size,
                   padding=self.padding, stride=self.stride)

        x = self.proj(x.permute(0, 2, 3, 1))
        x = self.proj_drop(x)

        return x



if __name__ == "__main__":
    # 如果GPU可用，将模块移动到 GPU
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # 输入张量 (batch_size, height, width,channels)
    x = torch.randn(1,40,40,32).to(device)
    # 初始化 OutlookAttention 模块
    dim=32
    block = OutlookAttention(dim,8)
    print(block)
    block = block.to(device)
    # 前向传播
    output = block(x)
    print("输入:", x.shape)
    print("输出:", output.shape)