YOLO算法改进Backbone系列之：EfficientViT

EfficientViT: Memory Effificient Vision Transformer with Cascaded Group Attention
摘要：视觉transformer由于其高模型能力而取得了巨大的成功。然而，它们卓越的性能伴随着沉重的计算成本，这使得它们不适合实时应用。在这篇论文中，本文提出了一个高速视觉transformer家族，名为EfficientViT。本文发现现有的transformer模型的速度通常受到内存低效操作的限制，特别是在MHSA中的张量重塑和单元函数。因此，本文设计了一种具有三明治布局的新构建块，即在高效FFN层之间使用单个内存绑定的MHSA，从而提高了内存效率，同时增强了信道通信。此外，本文发现注意图在头部之间具有很高的相似性，从而导致计算冗余。为了解决这个问题，本文提出了一个级联的群体注意模块，以不同的完整特征分割来馈送注意头，不仅节省了计算成本，而且提高了注意多样性。综合实验表明，高效vit优于现有的高效模型，在速度和精度之间取得了良好的平衡。例如，本文的EfficientViT-M5在准确率上比MobileNetV3-Large高出1.9%，而在Nvidia V100 GPU和Intel Xeon CPU上的吞吐量分别高出40.4%和45.2%。与最近的高效型号MobileViT-XXS相比，efficientvitt - m2的精度提高了1.8%，同时在GPU/CPU上运行速度提高了5.8 ×/3.7 ×，转换为ONNX格式时速度提高了7.4×

本文通过分析DeiT和Swin两个Transformer架构得出如下结论：

• 适当降低MHSA层利用率可以在提高模型性能的同时提高访存效率
• 在不同的头部使用不同的通道划分特征，而不是像MHSA那样对所有头部使用相同的全特征，可以有效地减少注意力计算冗余
• 典型的通道配置，即在每个阶段之后将通道数加倍或对所有块使用等效通道，可能在最后几个块中产生大量冗余
• 在维度相同的情况下，Q、K的冗余度比V大得多 a new building block with a sandwich
  layout（减少self-attention的次数）：之前是一个block self-attention->fc->self-attention->fc->self-attention->fc->…N次数；现在是一个blockfc->self-attention->fc；不仅能够提升内存效率而且能够增强通道间的计算
  cascaded group attention：让多头串联学习特征：第一个头学习完特征后，第二个头利用第一个头学习到的特征的基础上再去学习(原来的transformer是第二个头跟第一个头同时独立地去学习)，同理第三个头学习时也得利用上第二个头学习的结果再去学习

Efficientvit模型结构如下图所示：

a memory-efficient sandwich layout

Cascaded Group Attention：解决了原来模型中多头重复学习(学习到的特征很多都是相似的)的问题，这里每个头学到的特征都不同，而且越往下面的头学到的特征越丰富。

Q是主动查询的行为，特征比K更加丰富，所以额外做了个Token Interation
Q进行self-attention之前先通过多次分组卷积再一次学习
Parameter Reallocation
self-attention主要在进行QK，而且还需要对Q/K进行reshape，所以为了运算效率更快，Q与K的维度小一点
而V只在后面被QK得到的结果进行权重分配，没那么费劲，为了学习更多的特征，所以V维度更大一些

Efficientvit变体模型结构如下表所示：

在YOLOv5项目中添加EfficientViT模型作为Backbone使用的教程：
（1）将YOLOv5项目的models/yolo.py修改parse_model函数以及BaseModel的_forward_once函数

def parse_model(d, ch):  # model_dict, input_channels(3)
    # Parse a YOLOv5 model.yaml dictionary
    LOGGER.info(f"\n{
     '':>3}{
     'from':>18}{
     'n':>3}{
     'params':>10}  {
     'module':<40}{
     'arguments':<30}")
    anchors, nc, gd, gw, act = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple'], d.get('activation')
    if act:
        Conv.default_act = eval(act)  # redefine default activation, i.e. Conv.default_act = nn.SiLU()
        LOGGER.info(f"{
     colorstr('activation:')} {
     act}")  # print
    na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors  # number of anchors
    no = na * (nc + 5)  # number of outputs = anchors * (classes + 5)

    # ---------------------------------------------------------------------------------------------------
    is_backbone = False
    layers, save, c2 = [], [], ch[-1]  # layers, savelist, ch out
    for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']):  # from, number, module, args
        try:
            t = m
            m = eval(m) if isinstance(m, str) else m  # eval strings
        except:
            pass
        for j, a in enumerate(args):
            with contextlib.suppress(NameError):
                try:
                    args[j] = eval(a) if isinstance(a, str) else a  # eval strings
                except:
                    args[j] = a

        n = n_ = max(round(n * gd), 1) if n > 1 else n  # depth gain
        if m in {
   
                Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv,
                BottleneckCSP, C3, C3TR, C3SPP, C3Ghost, nn.ConvTranspose2d, DWConvTranspose2d, C3x}:
            c1, c2 = ch[f], args[0]
            if c2 != no:  # if not output
                c2 = make_divisible(c2 * gw, 8)

            args = [c1, c2, *args[1:]]
            if m in {
   BottleneckCSP, C3, C3TR, C3Ghost, C3x}:
                args.insert(2, n)  # number of repeats
                n = 1
        elif m is nn.BatchNorm2d:
            args = [ch[f]]
        elif m is Concat:
            c2 = sum(ch[x] for x in f)
        # TODO: channel, gw, gd
        elif m in {
   Detect, Segment}:
            args.append([ch[x] for x in f])
            if isinstance(args[1], int):  # number of anchors
                args[1] = [list(range(args[1] * 2))] * len(f)
            if m is Segment:
                args[3] = make_divisible(args[3] * gw, 8)
        elif m is Contract:
            c2 = ch[f] * args[0] ** 2
        elif m is Expand:
            c2 = ch[f] // args[0] ** 2
        # -------------------------------------------------------------------------------------
        elif m in {
   }:
            m = m(*args)
            c2 = m.channel
        # -------------------------------------------------------------------------------------
        else:
            c2 = ch[f]

        # -------------------------------------------------------------------------------------
        if isinstance(c2, list):
            is_backbone = True
            m_ = m
            m_.backbone = True
        else:
            m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args)  # module
            t = str(m)[8:-2].replace('__main__.', '')  # module type
        # -------------------------------------------------------------------------------------

        np = sum(x.numel() for x in m_.parameters())  # number params
        # -------------------------------------------------------------------------------------
        # m_.i, m_.f, m_.type, m_.np = i, f, t, np  # attach index, 'from' index, type, number params
        m_.i, m_.f, m_.type, m_.np = i + 4 if is_backbone else i, f, t, np  # attach index, 'from' index, type, number params
        # -------------------------------------------------------------------------------------
        
        LOGGER.info(f'{
     i:>3}{
     str(f):>18}{
     n_:>3}{
     np:10.0f}  {
     t:<40}{
     str(args):<30}')  # print
        save.extend(x % (i + 4 if is_backbone else i) for x in ([f] if isinstance(f, int) else f) if x != -1)  # append to savelist
        # save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1)  # append to savelist
        layers.append(m_)
        if i == 0:
            ch = []
        
        # -------------------------------------------------------------------------------------
        if isinstance(c2, list):
            ch.extend(c2)
            for _ in range(5 - len(ch)):
                ch.insert(0, 0)
        else:
            ch.append(c2)
        # -------------------------------------------------------------------------------------

    return nn.Sequential(*layers), sorted(save)


def _forward_once(self, x, profile=False, visualize=False):
        y, dt = [], []  # outputs
        for m in self.model:
            if m.f != -1:  # if not from previous layer
                x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f]  # from earlier layers
            if profile:
                self._profile_one_layer(m, x, dt)
            if hasattr(m, 'backbone'):
                x = m(x)
                for _ in range(5 - len(x)):
                    x.insert(0, None)
                for i_idx, i in enumerate(x):
                    if i_idx in self.save:
                        y.append(i)
                    else:
                        y.append(None)
                x = x[-1]
            else:
                x = m(x)  # run
                y.append(x if m.i in self.save else None)  # save output
            if visualize:
                feature_visualization(x, m.type, m.i, save_dir=visualize)
        return x

（2）在models/backbone（新建）文件下新建EfficientViT.py，添加如下的代码：

# --------------------------------------------------------