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卷积以及卷积。具体来说,我们将卷积拆分为depth-wise dilation卷积、depth-wise卷积以及卷积。通过上述分解,我们能够以少量的计算量与参数捕获长距离相关性。在得到长距离相关性后,我们可以估计点的重要性并生成注意力图。
上图a给出了所提LKA模块的示意图及其与其他模块的区别,它可以描述如下:
上表对比了卷积、自注意力以及LKA之间的区别,可以看到:LKA同时具有卷积与自注意力的特性LKA不仅具有空域自适应性,同时具有通道维度上的自适应性
classMlp(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)


defforward(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


classAttentionModule(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
)


defforward(self, x):
        u = x.clone()        

        attn = self.conv0(x)

        attn = self.conv_spatial(attn)

        attn = self.conv1(attn)

return
 u * attn



classSpatialAttention(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
)


defforward(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


classBlock(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
)


defforward(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采用深度卷积、扩展因子为3的深度卷积以及卷积近似卷积。基于该配置,VAN可以有效的获取局部信息以及长距离相关性。此外,VAN采用卷积进行初始的4倍下采样,然后卷积进行2倍下采样。
classOverlapPatchEmbed(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)


defforward(self, x):
        x = self.proj(x)

        _, _, H, W = x.shape

        x = self.norm(x)        

return
 x, H, W

3Experiments

4推荐阅读

  1. Timm助力ResNet焕发“第二春”,无蒸馏且无额外数据,性能高达80.4%
  2. PPLCNet:CPU端强悍担当,吊打现有主流轻量型网络,百度提出CPU端的最强轻量型架构
  3. GPU端精度最高速度最快的强悍担当:GENet
  4. 新一代移动端模型MobileNeXt来了!打破常规,逆残差模块超强改进,精度速度双超MobileNetV2
  5. 新坑!谷歌提出MLP-Mixer:一种无卷积、无注意力,纯MLP构成的视觉架构
继续阅读
阅读原文