EfficientFormerV2：重新审视ViT能否达到MobileNet的规模和速度

向岸看

已于 2023-10-31 11:50:09 修改

阅读量650

点赞数 4

分类专栏： ICCV会议文章标签： 1024程序员节 ICCV会议轻量ViT

于 2023-10-24 22:48:31 首次发布

本文链接：https://blog.csdn.net/qq_45981086/article/details/134021379

版权

ICCV会议专栏收录该内容

4 篇文章 0 订阅

订阅专栏

期刊：2023 IEEE/CVF International Conference on Computer Vision (ICCV)

标题：Rethinking（重新审视） Vision Transformers（ViT） for MobileNet Size and Speed（MobileNet的规模和速度）

——重新审视ViT能否达到MobileNet的规模和速度

论文：https://arxiv.org/pdf/2212.08059.pdf

源码：GitHub - snap-research/EfficientFormer: EfficientFormerV2 [ICCV 2023] & EfficientFormer [NeurIPs 2022]

一、摘要

二、网络架构

3.1. Token Mixers vs. Feed Forward Network（改进1：FFN 代替Token Mixers）

3.2 MHSA Improvements（改进2：改进多头注意力）

3.3 Attention on Higher Resolution（改进3：更高分辨率上的注意力）

3.4 Dual-Path Attention Downsampling（改进4：双路径注意力下采样）

三、EfficientFormerV2概述

一、摘要

研究背景：随着视觉Transformers（ViTs）在计算机视觉任务中的成功，最近的技术试图优化ViT的性能和复杂性，以实现在移动设备上的高效部署。研究人员提出了多种方法来加速注意力机制，改进低效设计，或结合mobile-friendly的轻量级卷积来形成混合架构。

研究问题：然而，ViT及其变体仍然比轻量级的CNNs具有更高的延迟或更多的参数，即使对于多年前的MobileNet也是如此。实际上，延迟和大小对于资源受限硬件上的高效部署都至关重要。

主要工作：

1. 重新审视了ViT的设计选择，并提出了一种具有低延迟和高参数效率的改进型超网络。
2. 进一步引入了一种细粒度联合搜索策略，该策略可以通过同时优化延迟和参数量来找到有效的架构。

研究成果：所提出的模型EfficientFormerV2在ImageNet-1K上实现了比MobileNetV2和MobileNetV1高约4%的top-1精度，具有相似的延迟和参数。论文证明，适当设计和优化的ViT可以以MobileNet级别的大小和速度实现高性能。

二、网络架构

EfficientFormerV2网络架构如下：

(a)Efficientformer基线 (b)统一的FFN (c)MHSA的改进 (d)&(e)注意更高的分辨率 (f)双路径注意下降采样

注意：图b -> Local模块、图c与图f（图f中的attention即为图c） -> Global模块，才是真正的EfficientFormerV2的网络架构。其余图像皆为改进过程。

以下是三个改进方法。

3.1. Token Mixers vs. Feed Forward Network（改进1：FFN 代替Token Mixers）

前人工作：PoolFormer和EfficientFormer采用3 x 3个平均池化层作为local token(局部标记)混合器。（local token混合器的作用）增强局部信息可以提高性能，并使ViT在没有显式位置嵌入的情况下更加鲁棒。

动机1：用相同内核大小的深度卷积（DWCONV）替换这些层不会引入延迟开销，而性能提高了0.6%，额外参数可以忽略不计（0.02M）。

动机2：此外，最近的工作表明，在ViT中的前馈网络（FFN）中添加局部信息捕获层也是有益的，以提高性能，同时开销较小。

改进：基于这些观察，删除了显式残差连接的局部标记混合器，并将深度方向的 3 × 3 CONV 移动到FFN中，以获得启用局部性的统一FFN。将统一的 FFN 应用于网络的所有阶段。这种设计修改将网络架构简化为只有两种类型的块 (Local FFN 和 Gobel Attention)，并在相同延迟下将准确率提高到 80.3%，参数开销很小(0.1M)。

代码如下：

FFN深度卷积前馈网络（Local 模块）

# 深度卷积前馈网络
class FFN(nn.Module):
    def __init__(self, dim, pool_size=3, mlp_ratio=4.,
                 act_layer=nn.GELU,
                 drop=0., drop_path=0.,
                 use_layer_scale=True, layer_scale_init_value=1e-5):
        super().__init__()

        mlp_hidden_dim = int(dim * mlp_ratio)       # 隐藏层维度
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim,
                       act_layer=act_layer, drop=drop, mid_conv=True)       # 多层感知机

        self.drop_path = DropPath(drop_path) if drop_path > 0. \
            else nn.Identity()
        self.use_layer_scale = use_layer_scale
        if use_layer_scale:
            self.layer_scale_2 = nn.Parameter(
                layer_scale_init_value * torch.ones(dim).unsqueeze(-1).unsqueeze(-1), requires_grad=True)   # 缩放因子

    def forward(self, x):
        if self.use_layer_scale:
            x = x + self.drop_path(self.layer_scale_2 * self.mlp(x))    # 残差连接
        else:
            x = x + self.drop_path(self.mlp(x))
        return x

Mlp多层感知机

# 多层感知机
class Mlp(nn.Module):
    """
    Implementation of MLP with 1*1 convolutions.
    Input: tensor with shape [B, C, H, W]
    """

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

        if self.mid_conv:
            self.mid = nn.Conv2d(hidden_features, hidden_features, kernel_size=3, stride=1, padding=1,
                                 groups=hidden_features)
            self.mid_norm = nn.BatchNorm2d(hidden_features)

        self.norm1 = nn.BatchNorm2d(hidden_features)
        self.norm2 = nn.BatchNorm2d(out_features)

    def _init_weights(self, m):
        if isinstance(m, nn.Conv2d):
            trunc_normal_(m.weight, std=.02)
            if m.bias is not None:
                nn.init.constant_(m.bias, 0)

    def forward(self, x):
        x = self.fc1(x)     # 1x1卷积
        x = self.norm1(x)
        x = self.act(x)     # 激活层，GELU激活

        if self.mid_conv:
            x_mid = self.mid(x)     # 3x3卷积
            x_mid = self.mid_norm(x_mid)
            x = self.act(x_mid)
        x = self.drop(x)

        x = self.fc2(x)     # 1x1卷积
        x = self.norm2(x)

        x = self.drop(x)
        return x

3.2 MHSA Improvements（改进2：改进多头注意力）

前人工作：例如，最早的Vision Transformer(ViT)和Efficientformer的注意力模块都使用了MHSA + FFN的结构。

动机：为了在不增加模型大小和延迟的额外开销的情况下提高注意模块性能。

改进MHSA的两种方法：

1. 首先，通过添加深度方向的3 × 3 CONV将局部信息注入到值矩阵（V）中。
2. 其次，通过在头部维度上添加全连接的层（talking head）来实现注意力头部之间的通信。

代码如下：

多头自注意力MHSA

class Attention4D(torch.nn.Module):
    def __init__(self, dim=384, key_dim=32, num_heads=8,
                 attn_ratio=4,
                 resolution=7,
                 act_layer=nn.ReLU,
                 stride=None):
        super().__init__()
        self.num_heads = num_heads
        self.scale = key_dim ** -0.5
        self.key_dim = key_dim
        self.nh_kd = nh_kd = key_dim * num_heads

        if stride is not None:
            self.resolution = math.ceil(resolution / stride)
            self.stride_conv = nn.Sequential(nn.Conv2d(dim, dim, kernel_size=3, stride=stride, padding=1, groups=dim),
                                             nn.BatchNorm2d(dim), )
            self.upsample = nn.Upsample(scale_factor=stride, mode='bilinear')
        else:
            self.resolution = resolution
            self.stride_conv = None
            self.upsample = None

        self.N = self.resolution ** 2
        self.N2 = self.N
        self.d = int(attn_ratio * key_dim)
        self.dh = int(attn_ratio * key_dim) * num_heads
        self.attn_ratio = attn_ratio
        h = self.dh + nh_kd * 2
        self.q = nn.Sequential(nn.Conv2d(dim, self.num_heads * self.key_dim, 1),
                               nn.BatchNorm2d(self.num_heads * self.key_dim), )
        self.k = nn.Sequential(nn.Conv2d(dim, self.num_heads * self.key_dim, 1),
                               nn.BatchNorm2d(self.num_heads * self.key_dim), )
        self.v = nn.Sequential(nn.Conv2d(dim, self.num_heads * self.d, 1),
                               nn.BatchNorm2d(self.num_heads * self.d),
                               )
        self.v_local = nn.Sequential(nn.Conv2d(self.num_heads * self.d, self.num_heads * self.d,
                                               kernel_size=3, stride=1, padding=1, groups=self.num_heads * self.d),
                                     nn.BatchNorm2d(self.num_heads * self.d), )
        self.talking_head1 = nn.Conv2d(self.num_heads, self.num_heads, kernel_size=1, stride=1, padding=0)
        self.talking_head2 = nn.Conv2d(self.num_heads, self.num_heads, kernel_size=1, stride=1, padding=0)

        self.proj = nn.Sequential(act_layer(),
                                  nn.Conv2d(self.dh, dim, 1),
                                  nn.BatchNorm2d(dim), )

        points = list(itertools.product(range(self.resolution), range(self.resolution)))
        N = len(points)
        attention_offsets = {}
        idxs = []
        for p1 in points:
            for p2 in points:
                offset = (abs(p1[0] - p2[0]), abs(p1[1] - p2[1]))
                if offset not in attention_offsets:
                    attention_offsets[offset] = len(attention_offsets)
                idxs.append(attention_offsets[offset])
        self.attention_biases = torch.nn.Parameter(
            torch.zeros(num_heads, len(attention_offsets)))
        self.register_buffer('attention_bias_idxs',
                             torch.LongTensor(idxs).view(N, N))

    @torch.no_grad()
    def train(self, mode=True):
        super().train(mode)
        if mode and hasattr(self, 'ab'):
            del self.ab
        else:
            self.ab = self.attention_biases[:, self.attention_bias_idxs]

    def forward(self, x):  # x (B,N,C)
        B, C, H, W = x.shape
        if self.stride_conv is not None:
            x = self.stride_conv(x)

        q = self.q(x).flatten(2).reshape(B, self.num_heads, -1, self.N).permute(0, 1, 3, 2)     # 降维为2维，转置
        k = self.k(x).flatten(2).reshape(B, self.num_heads, -1, self.N).permute(0, 1, 2, 3)
        v = self.v(x)
        v_local = self.v_local(v)       # 通过一个3x3的卷积将局部信息注入V中
        v = v.flatten(2).reshape(B, self.num_heads, -1, self.N).permute(0, 1, 3, 2)     # 降维为2维，转置

        attn = (
                (q @ k) * self.scale    # 矩阵乘法，计算相似度
                +
                (self.attention_biases[:, self.attention_bias_idxs]     # 融合一个位置编码
                 if self.training else self.ab)
        )
        # attn = (q @ k) * self.scale
        attn = self.talking_head1(attn)     # 1x1卷积->全连接，实现注意力头部之间的通信
        attn = attn.softmax(dim=-1)
        attn = self.talking_head2(attn)     # 同上

        x = (attn @ v)      # 注意力融合

        out = x.transpose(2, 3).reshape(B, self.dh, self.resolution, self.resolution) + v_local     # 最后再与v_local融合
        if self.upsample is not None:
            out = self.upsample(out)

        out = self.proj(out)    # 输出再进行激活 + 卷积 + 正则
        return out

AttnFFN（Local Global 模块）

class AttnFFN(nn.Module):
    def __init__(self, dim, mlp_ratio=4.,
                 act_layer=nn.ReLU, norm_layer=nn.LayerNorm,
                 drop=0., drop_path=0.,
                 use_layer_scale=True, layer_scale_init_value=1e-5,
                 resolution=7, stride=None):

        super().__init__()

        self.token_mixer = Attention4D(dim, resolution=resolution, act_layer=act_layer, stride=stride)      # MHSA多头自注意力
        mlp_hidden_dim = int(dim * mlp_ratio)       # 隐藏层维度
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim,
                       act_layer=act_layer, drop=drop, mid_conv=True)       # 深度卷积

        self.drop_path = DropPath(drop_path) if drop_path > 0. \
            else nn.Identity()      # drop_path概率
        self.use_layer_scale = use_layer_scale
        if use_layer_scale:     # 缩放因子
            self.layer_scale_1 = nn.Parameter(
                layer_scale_init_value * torch.ones(dim).unsqueeze(-1).unsqueeze(-1), requires_grad=True)
            self.layer_scale_2 = nn.Parameter(
                layer_scale_init_value * torch.ones(dim).unsqueeze(-1).unsqueeze(-1), requires_grad=True)

    def forward(self, x):
        if self.use_layer_scale:
            x = x + self.drop_path(self.layer_scale_1 * self.token_mixer(x))    # 多头自注意力 + 残差连接
            x = x + self.drop_path(self.layer_scale_2 * self.mlp(x))    # mlp深度卷积 + 残差连接

        else:
            x = x + self.drop_path(self.token_mixer(x))
            x = x + self.drop_path(self.mlp(x))
        return x

3.3 Attention on Higher Resolution（改进3：更高分辨率上的注意力）

研究问题：注意力机制有利于提高效果。然而，将其应用于高分辨率特征损害了移动端的效率，因为它具有与空间分辨率对应的二次时间复杂度。

动机：能不能通过降低输入分辨率来提高注意力机制的效率？

前人工作：虽然一些工作提出了基于窗口的注意力或下采样的键和值来缓解这个问题，但我们发现它们不是移动的部署的最佳选择。基于窗口的注意力很难在移动的设备上加速，因为复杂的窗口分区和重新排序（作者放弃了这种方法）。对于下采样键（K）和值（V），需要全分辨率查询（Q）以在注意矩阵乘法之后保持输出分辨率（Out）（作者最终采用）。

改进方法：将所有查询、关键字和值下采样到固定的空间分辨率（1/32），并将来自注意力的输出内插回到原始分辨率以馈送到下一层，作者将这种方法称为“跨步注意力”。

3.4 Dual-Path Attention Downsampling（改进4：双路径注意力下采样）

前人工作：LeViT 和 UniNet 提出通过注意机制将特征分辨率减半，具体来说，查询中的token数量减少了一半，从而使注意模块的输出被降采样。（作者不赞同这种做法，因为Query的token数量是很重要的）

改进方案：提出了一种组合策略，即双路径注意力下采样，它既利用了局部性又利用了全局依赖性。

1. 使用池化作为静态局部下采样。
2. 3 × 3 DWCONV作为可学习的局部下采样。

并将结果合并投影到查询中。

代码如下：

LGQuery（双路径注意力下采样）

class LGQuery(torch.nn.Module):
    def __init__(self, in_dim, out_dim, resolution1, resolution2):
        super().__init__()
        self.resolution1 = resolution1
        self.resolution2 = resolution2
        self.pool = nn.AvgPool2d(1, 2, 0)
        self.local = nn.Sequential(nn.Conv2d(in_dim, in_dim, kernel_size=3, stride=2, padding=1, groups=in_dim),
                                   )
        self.proj = nn.Sequential(nn.Conv2d(in_dim, out_dim, 1),
                                  nn.BatchNorm2d(out_dim), )

    def forward(self, x):
        B, C, H, W = x.shape
        local_q = self.local(x)
        pool_q = self.pool(x)
        q = local_q + pool_q            # 双路径下采
        q = self.proj(q)
        return q

Attention模块（高分辨版）

class Attention4DDownsample(torch.nn.Module):
    def __init__(self, dim=384, key_dim=16, num_heads=8,
                 attn_ratio=4,
                 resolution=7,
                 out_dim=None,
                 act_layer=None,
                 ):
        super().__init__()
        self.num_heads = num_heads
        self.scale = key_dim ** -0.5
        self.key_dim = key_dim
        self.nh_kd = nh_kd = key_dim * num_heads

        self.resolution = resolution

        self.d = int(attn_ratio * key_dim)
        self.dh = int(attn_ratio * key_dim) * num_heads
        self.attn_ratio = attn_ratio
        h = self.dh + nh_kd * 2

        if out_dim is not None:
            self.out_dim = out_dim
        else:
            self.out_dim = dim

        self.resolution2 = math.ceil(self.resolution / 2)
        self.q = LGQuery(dim, self.num_heads * self.key_dim, self.resolution, self.resolution2)         # 双注意力下采样

        self.N = self.resolution ** 2
        self.N2 = self.resolution2 ** 2

        self.k = nn.Sequential(nn.Conv2d(dim, self.num_heads * self.key_dim, 1),
                               nn.BatchNorm2d(self.num_heads * self.key_dim), )
        self.v = nn.Sequential(nn.Conv2d(dim, self.num_heads * self.d, 1),
                               nn.BatchNorm2d(self.num_heads * self.d),
                               )
        self.v_local = nn.Sequential(nn.Conv2d(self.num_heads * self.d, self.num_heads * self.d,
                                               kernel_size=3, stride=2, padding=1, groups=self.num_heads * self.d),
                                     nn.BatchNorm2d(self.num_heads * self.d), )

        self.proj = nn.Sequential(
            act_layer(),
            nn.Conv2d(self.dh, self.out_dim, 1),
            nn.BatchNorm2d(self.out_dim), )

        points = list(itertools.product(range(self.resolution), range(self.resolution)))
        points_ = list(itertools.product(
            range(self.resolution2), range(self.resolution2)))
        N = len(points)
        N_ = len(points_)
        attention_offsets = {}
        idxs = []
        for p1 in points_:
            for p2 in points:
                size = 1
                offset = (
                    abs(p1[0] * math.ceil(self.resolution / self.resolution2) - p2[0] + (size - 1) / 2),
                    abs(p1[1] * math.ceil(self.resolution / self.resolution2) - p2[1] + (size - 1) / 2))
                if offset not in attention_offsets:
                    attention_offsets[offset] = len(attention_offsets)
                idxs.append(attention_offsets[offset])

        self.register_buffer('attention_biases', torch.zeros(num_heads, 196))
        self.register_buffer('attention_bias_idxs',
                             torch.ones(49, 196).long())

        self.attention_biases_seg = torch.nn.Parameter(
            torch.zeros(num_heads, len(attention_offsets)))
        self.register_buffer('attention_bias_idxs_seg',
                             torch.LongTensor(idxs).view(N_, N))

    @torch.no_grad()
    def train(self, mode=True):
        super().train(mode)
        if mode and hasattr(self, 'ab'):
            del self.ab
        else:
            self.ab = self.attention_biases_seg[:, self.attention_bias_idxs_seg]

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

        q = self.q(x).flatten(2).reshape(B, self.num_heads, -1, H * W // 4).permute(0, 1, 3, 2)
        k = self.k(x).flatten(2).reshape(B, self.num_heads, -1, H * W).permute(0, 1, 2, 3)
        v = self.v(x)
        v_local = self.v_local(v)
        v = v.flatten(2).reshape(B, self.num_heads, -1, H * W).permute(0, 1, 3, 2)

        attn = (q @ k) * self.scale                     # (H * W // 4, H * W)
        bias = self.attention_biases_seg[:, self.attention_bias_idxs_seg] if self.training else self.ab
        bias = torch.nn.functional.interpolate(bias.unsqueeze(0), size=(attn.size(-2), attn.size(-1)), mode='bicubic')
        attn = attn + bias

        attn = attn.softmax(dim=-1)

        x = (attn @ v).transpose(2, 3)                  # (H * W // 4, H * W)
        out = x.reshape(B, self.dh, H // 2, W // 2) + v_local

        out = self.proj(out)
        return out

三、EfficientFormerV2概述

EfficientFormerV2的重要模块定义如下：

1.嵌入式输入

B 表示批大小，C 表示通道尺寸(也表示网络的宽度)，H 和 W 为特征的高度和宽度， $X_j$ 为第 j 阶段的特征，j∈{1, 2, 3, 4} 和 i 表示第 i 层。

2.Local模块：FFN

其中， $S_{i, j}$ 是可学习的层缩放因子，阶段宽度 $C_j$ 和每个块的扩展比 $E_{i,j}$ 。

3.Gobal模块：MHSA + FFN

其中，查询(Q)、Keys (K)和值(V)通过是线性层Q、K、V得到的输入特征。

4.MHSA

其中，ab是位置编码的可学习注意偏差。

四、实验

4.1 单指标评估

分别针对模型大小、大模型、推理速度和联合优化方面进行对比实验。

对于模型大小，EfficientFormerV 2-S0比EdgeViT-XXS超出了1.3%的top-1精度，甚至少了0.6M参数，比MobileNetV 2 ×1.0优于3.5%的top-1，参数数量相似。对于大型模型，EfficientFormerV 2-L模型实现了与最近的EfficientFormerL 7 相同的精度，同时小3.1倍。在速度方面，在延迟相当或更低的情况下，EfficientFormerV2-S2的性能分别优于UniNet-B1，EdgeViT-S和EfficientFormerL 1，分别为0.8%，0.6%和2.4%。 EiffcientFormer V2-S1的效率分别比MobileViT-XS、EdgeViT-XXS 和EdgeViTXS 高出4.2%、4.6%和1.5%，其中MES要高得多。

4.2 下游任务

对象检测和实例分割

对象检测和实例分割。具有相似的模型尺寸，EfficientFormerV2-S2性能比PoolFormer S12高出 6.1 APbox和 4.9 APmask。EfficientFormerV2-L的效率EfficientFormer-L3高出3.3 APbox和2.3 APmask。语义分割。EfficientFormerV2-S2的效率分别比Poolformer-s12和Efficientformer-L1高5.2和3.5 mIoU。

4.3 搜索算法中的消融分析

五、总结

1. 本文重心在于搜索策略，提出了一个关于大小和速度的细粒度联合搜索，以及通过这个搜索策略得到了最佳的Efficientformer架构。

2. 为了在更高分辨率的图像上实现注意力，作者引入双路径注意力下采样的机制，使用池化作为静态局部下采样，3 × 3 DWCONV作为可学习的局部下采样。

3. EfficientFormerV2是轻量级的，超快的推理速度和高性能。

向岸看

关注

4
点赞
踩
6

收藏

觉得还不错? 一键收藏
打赏
3
评论
EfficientFormerV2：重新审视ViT能否达到MobileNet的规模和速度

Rethinking Vision Transformers for MobileNet Size and Speed论文总结和代码分析。
复制链接

扫一扫