[人工智能]EdgeNeXt: Efficiently Amalgamated CNN-Transformer Architecture for Mobile Vision Applications

摘要

众所周知，模型的准确性很大程度上可以评判一个模型的优劣。为了提高模型的准确性，人们往往会设计复杂的神经网络来提高准确率。然而，模型越复杂，需要的计算资源就越大，这就导致了大模型无法部署在例如手机等边缘设备上。本文一方面是为了结合CNN和Transformer的优势，一方面是为了使得模型能在边缘设备上部署，提出了EdgeNext。

介绍

CNN和Transformer的对比

CNN具备局部感受野，无法对全局信息进行建模。
CNN学习到的权重在推理过程中是静态的，无法灵活的适应输入的内容。
Transformer可以缓解这个问题，但是计算的复杂度太高，对边缘设备不友好。

贡献

提出了SDTA（Split depth-wise transpose attention），可以有效的增加全局和局部表示。最主要的是不增加参数和乘法操作。

相关工作

总体架构

总体架构

Convolution Encoder

类似于深度可分离卷积，使用每个阶段大小可变的深度卷积来丰富局部表示。
然后使用1*1的卷积进行特征组合，进行不同通道上特征信息的交互。
空间混合（深度卷积）和通道混合（点卷积）

SDTA Encoder

split通过编码输入图像中的各种空间级别来学习自适应多尺度特征表示。
自注意力的计算在通道维度上进行，大大减少了计算的复杂度。

代码解析

EdgeNeXt有两种实现。一种是LayerNorm和GELU的，一种为Hard-Swish和BatchNorm。主要以介绍LayerNorm和GELU的ConvEncoder和SDTA。

ConvEncoder

class ConvEncoder(nn.Module):
    def __init__(self, dim, drop_path=0., layer_scale_init_value=1e-6, expan_ratio=4, kernel_size=7):
        super().__init__()
        # 空间混合
        self.dwconv = nn.Conv2d(dim, dim, kernel_size=kernel_size, padding=kernel_size // 2, groups=dim)
        self.norm = LayerNorm(dim, eps=1e-6)
        # 通道混合
        self.pwconv1 = nn.Linear(dim, expan_ratio * dim)
        self.act = nn.GELU()
        self.pwconv2 = nn.Linear(expan_ratio * dim, dim)
        self.gamma = nn.Parameter(layer_scale_init_value * torch.ones(dim),
                                  requires_grad=True) if layer_scale_init_value > 0 else None
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()

    def forward(self, x):
        input = x
        x = self.dwconv(x)
        x = x.permute(0, 2, 3, 1)  # (N, C, H, W) -> (N, H, W, C)
        x = self.norm(x)
        x = self.pwconv1(x)
        x = self.act(x)
        x = self.pwconv2(x)
        if self.gamma is not None:
            x = self.gamma * x
        x = x.permute(0, 3, 1, 2)  # (N, H, W, C) -> (N, C, H, W)

        x = input + self.drop_path(x)
        return x

SDTA(Split Depth-wise Transpose Attention)

class SDTAEncoder(nn.Module):
    def __init__(self, dim, drop_path=0., layer_scale_init_value=1e-6, expan_ratio=4,
                 use_pos_emb=True, num_heads=8, qkv_bias=True, attn_drop=0., drop=0., scales=1):
        super().__init__()
        width = max(int(math.ceil(dim / scales)), int(math.floor(dim // scales)))
        self.width = width
        if scales == 1:
            self.nums = 1
        else:
            self.nums = scales - 1
        convs = []
        for i in range(self.nums):
            convs.append(nn.Conv2d(width, width, kernel_size=3, padding=1, groups=width))
        self.convs = nn.ModuleList(convs)

        self.pos_embd = None
        if use_pos_emb:
            self.pos_embd = PositionalEncodingFourier(dim=dim)
        self.norm_xca = LayerNorm(dim, eps=1e-6)
        self.gamma_xca = nn.Parameter(layer_scale_init_value * torch.ones(dim),
                                      requires_grad=True) if layer_scale_init_value > 0 else None
        self.xca = XCA(dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop)

        self.norm = LayerNorm(dim, eps=1e-6)
        self.pwconv1 = nn.Linear(dim, expan_ratio * dim)  # pointwise/1x1 convs, implemented with linear layers
        self.act = nn.GELU()  # TODO: MobileViT is using 'swish'
        self.pwconv2 = nn.Linear(expan_ratio * dim, dim)
        self.gamma = nn.Parameter(layer_scale_init_value * torch.ones((dim)),
                                  requires_grad=True) if layer_scale_init_value > 0 else None
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()

    def forward(self, x):
        input = x

        # 切分通道，加入深度卷积来增加多尺度表示
        spx = torch.split(x, self.width, 1)
        for i in range(self.nums):
            if i == 0:
                sp = spx[i]
            else:
                sp = sp + spx[i]
            sp = self.convs[i](sp)
            if i == 0:
                out = sp
            else:
                out = torch.cat((out, sp), 1)
        x = torch.cat((out, spx[self.nums]), 1)
        # XCA，在通道上计算注意力
        B, C, H, W = x.shape
        x = x.reshape(B, C, H * W).permute(0, 2, 1)
        if self.pos_embd:
            # 加入位置坐标
            pos_encoding = self.pos_embd(B, H, W).reshape(B, -1, x.shape[1]).permute(0, 2, 1)
            x = x + pos_encoding
        x = x + self.drop_path(self.gamma_xca * self.xca(self.norm_xca(x)))
        x = x.reshape(B, H, W, C)

        # Inverted Bottleneck，倒残差结构。
        # 倒残差结构的特点是：先对输入特征通道扩张，再提取特征，最后输出相应的特征通道，对于通道数来说有点中间大两头小类似于梭子的形状，所以称这样的结构为 Inverted residuals（倒残差结构）。
        x = self.norm(x)
        x = self.pwconv1(x)
        x = self.act(x)
        x = self.pwconv2(x)
        if self.gamma is not None:
            x = self.gamma * x
        x = x.permute(0, 3, 1, 2)  # (N, H, W, C) -> (N, C, H, W)

        x = input + self.drop_path(x)

        return x

github地址

edgeNext