优于ConvNeXt，南开&清华开源基于大核注意力的VAN架构

AIWalker

发布于 2022-03-29 17:04:50

9630

文章被收录于专栏：AIWalkerAIWalker

paper: https://arxiv.org/abs/2202.09741 code: https://github.com/Visual-Attention-Network

2022以来，CNN动静不小，先是有Meta(原Facebook)的ConvNeXt引发极大关注；近日南开程明明与其博士导师胡事民团队开源了基于全新大核注意力的VAN架构，一种99%纯度的新型CNN架构，在图像分类、目标检测、语义分割以及实例分割等诸多任务上均取得了超越Transformer、CNN的性能。推荐指数五颗星！

1出发点

尽管Transformer在CV领域引发轰动，但图像的2D特性导致自注意力面临如下三个挑战：

将图像视作1D序列将忽视其2D结构信息；
高分辨率图像会带来极大的计算复杂度；
它仅考虑了空域自适应，忽视了通道自适应性。

针对上述挑战，本文提出一种新的大核注意力(Large Kernel Attention，LKA)模块促进self-adatpive与long-range相关性，同时避免上述问题；与此同时，基于所提LKA构建了VAN架构，VAN在图像分类、目标检测、语义分割以及实例分割等任务上均取得了SOTA性能，超越了其他Transformer与CNN架构，包含最新的ConvXNet。

2Method

注意力可以视作一种自适应选择过程，它可以根据输入选择具有判别性的特征并自动忽视噪声响应。注意力的关键是步骤是生成表征不同点重要性的注意力图，故我们需要学习不同点之间的相关性。

有两种知名的方案可以构建不同点之间的相关性：自注意力与大核卷积。为克服两者的缺陷并利用其优势，我们对大核卷积进行分解以捕获长距离相关性。如上图所示，大核卷积可以拆分为三个成分：depth-wise卷积、depth-wise dilation卷积以及1\times 1 卷积。具体来说，我们将K \times K 卷积拆分为\frac{k}{d} \times \frac{k}{d} depth-wise dilation卷积、(2d-1) \times (2d-1) depth-wise卷积以及1\times 1 卷积。通过上述分解，我们能够以少量的计算量与参数捕获长距离相关性。在得到长距离相关性后，我们可以估计点的重要性并生成注意力图。

上图a给出了所提LKA模块的示意图及其与其他模块的区别，它可以描述如下：

Attention = Conv_{1\times 1}(DW-D-COnv(DW-Conv(F))) \\ Output = Attention \otimes F

上表对比了卷积、自注意力以及LKA之间的区别，可以看到：LKA同时具有卷积与自注意力的特性；LKA不仅具有空域自适应性，同时具有通道维度上的自适应性。

class Mlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Conv2d(in_features, hidden_features, 1)
        self.dwconv = DWConv(hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Conv2d(hidden_features, out_features, 1)
        self.drop = nn.Dropout(drop)
        self.apply(self._init_weights)

    def forward(self, x):
        x = self.fc1(x)
        x = self.dwconv(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x
        
class AttentionModule(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.conv0 = nn.Conv2d(dim, dim, 5, padding=2, groups=dim)
        self.conv_spatial = nn.Conv2d(dim, dim, 7, stride=1, padding=9, groups=dim, dilation=3)
        self.conv1 = nn.Conv2d(dim, dim, 1)

    def forward(self, x):
        u = x.clone()        
        attn = self.conv0(x)
        attn = self.conv_spatial(attn)
        attn = self.conv1(attn)
        return u * attn


class SpatialAttention(nn.Module):
    def __init__(self, d_model):
        super().__init__()

        self.proj_1 = nn.Conv2d(d_model, d_model, 1)
        self.activation = nn.GELU()
        self.spatial_gating_unit = AttentionModule(d_model)
        self.proj_2 = nn.Conv2d(d_model, d_model, 1)

    def forward(self, x):
        shorcut = x.clone()
        x = self.proj_1(x)
        x = self.activation(x)
        x = self.spatial_gating_unit(x)
        x = self.proj_2(x)
        x = x + shorcut
        return x
        
class Block(nn.Module):
    def __init__(self, dim, mlp_ratio=4., drop=0.,drop_path=0., act_layer=nn.GELU):
        super().__init__()
        self.norm1 = nn.BatchNorm2d(dim)
        self.attn = SpatialAttention(dim)
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()

        self.norm2 = nn.BatchNorm2d(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
        layer_scale_init_value = 1e-2            
        self.layer_scale_1 = nn.Parameter(
            layer_scale_init_value * torch.ones((dim)), requires_grad=True)
        self.layer_scale_2 = nn.Parameter(
            layer_scale_init_value * torch.ones((dim)), requires_grad=True)

    def forward(self, x):
        x = x + self.drop_path(self.layer_scale_1.unsqueeze(-1).unsqueeze(-1) * self.attn(self.norm1(x)))
        x = x + self.drop_path(self.layer_scale_2.unsqueeze(-1).unsqueeze(-1) * self.mlp(self.norm2(x)))
        return x

上表给出了基于LKA构建的VAN的架构参数配置信息，需要注意的是：虽然VAN在极力避免使用Transformer相关的算子，但仍用到了LN，故算不上“100%纯度”CNN😓。

在实现细节方面，LKA采用5\times 5 深度卷积、扩展因子为3的7\times 7 深度卷积以及1\times 1 卷积近似21\times 21 卷积。基于该配置，VAN可以有效的获取局部信息以及长距离相关性。此外，VAN采用7\times 7 卷积进行初始的4倍下采样，然后3\times 3 卷积进行2倍下采样。

class OverlapPatchEmbed(nn.Module):
    """ Image to Patch Embedding
    """

    def __init__(self, img_size=224, patch_size=7, stride=4, in_chans=3, embed_dim=768):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)

        self.img_size = img_size
        self.patch_size = patch_size
        self.H, self.W = img_size[0] // patch_size[0], img_size[1] // patch_size[1]
        self.num_patches = self.H * self.W
        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride,
                              padding=(patch_size[0] // 2, patch_size[1] // 2))
        self.norm = nn.BatchNorm2d(embed_dim)

    def forward(self, x):
        x = self.proj(x)
        _, _, H, W = x.shape
        x = self.norm(x)        
        return x, H, W