学习资源站

31-添加iRMB注意力机制（ICCV 2023_即插即用的反向残差注意力机制）

YOLOv5改进系列（30）——添加iRMB注意力机制（ICCV 2023|即插即用的反向残差注意力机制）

 

🚀 一、iRMB介绍 

 学习资料：

2023 腾讯优图/浙大/北大提出：重新思考高效神经模型的移动模块

1.1 iRMB简介

背景 

  • 基于Depth-Wise Convolution的IRB成为轻量化卷积模型的基础构成模块，但受限于静态CNN归纳偏置影响，这些模型的效果（尤其是下游检测、分割等任务）仍有待进一步提高。
  • 得益于Transformer的动态建模能力，基于ViT（Vison Transformer）的许多工作（如DeiT、PVT、SwinTransformer等）较CNN模型取得了明显的效果提升。然而，受限于MHSA（Multi-Head Self-Attention，多头自注意力）的二次方计算量需求，Transformer很少应用于移动端场景。

之前的尝试

因此一些研究人员尝试设计融合上述两者的混合模型，但一般会引入复杂的结构或多个混合模块，

本文的方法

这篇文章旨在设计轻量级高效的模型结构，平衡了参数量、计算量和精度.

本文首先抽象了MMB（Meta Mobile Block）用来对IRB和Transformer中的MHSA(Multi-Head Self-Attention，多头自注意力）/FFN(Feed-Forward Network，前向模型）进行归纳，其次实例化了高效的iRMB（Inverted Residual Mobile Block），最后仅使用该模块构建了高效的EMO（Efficient MOdel）轻量化主干模型。

实现效果

  • 1M/2M/5M 尺度下的 EMO 模型在 ImageNet-1K 的分类数据集上分别获得了 71.5/75.1/78.4 Top-1精度，超过了当前 SoTA 基于 CNN/Transformer 的轻量化模型；

  • 基于 SSDLite 框架仅使用 0.6G/0.9G/1.8G FLOPs 获得了 22.0/25.2/27.9 mAP；

  • 基于 DeepLabv3框架仅使用 2.4G/3.5G/5.8G FLOPs 获得了 33.5/35.3/37.8 mIoU，取了较好的 mobile 条件下的效果。

1.2 iRMB结构

iRMB结合了CNN的轻量级特性和Transformer模型的动态处理能力，这种结构特别适用于移动设备上需要进行密集预测任务的场景，因为它旨在在计算资源有限的情况下提供高效性能。

图2：左侧：抽象的统一元移动块，来源于多头自注意力/前馈网络和倒置残差块。归纳块可以通过不同的扩张比例λ和高效运算符F 推导成具体模块。右侧：类似ResNet的EMO，只由推导出的iRMB组成。

在图2中归纳出一个通用的Meta Mobile块（MMB），其中采用参数化参数扩张比例λ和高效运算符F 来实例化不同模块。

元移动模块

与MetaFormer的关系

基于Meta Mobile Block，设计了一个反向残差移动块 (iRMB)，它吸收了 CNN 架构的效率来建模局部特征和 Transformer 架构动态建模的能力来学习长距离交互。

如上图左侧所示，通过对 MobileNetv2 中的 IRB 以及 Transformer 中的核心 MHSA 和 FFN 模块进行抽象，本文提出了统一的 MMB 对上述几个结构进行归纳表示，即采用扩张率 𝜆 和高效算子 𝐹 来实例化不同的模块。

反向残差模块

反向残差模块（Inverted Residual Block，IRB）是一种用于深度神经网络中的模块结构。在iRMB添加倒置残差块（IRB）得模型能够更有效地处理长距离信息。

什么是倒残差模块？

通常情况下，残差连接会将输入添加到层的输出中，而在倒置残差块中，先将输入通过一个轻量级的层（比如1x1卷积）进行降维，然后再将降维后的特征进行处理，并最终将处理后的结果与原始输入相加，形成残差连接。这样的设计可以减少计算量和参数数量，同时仍然可以保持模型的性能和效率。

在iRMB中基于Meta Mobile Block设计一个反向残差移动块 (iRMB)，它吸收了 CNN 架构的效率来建模局部特征和 Transformer 架构动态建模的能力来学习长距离交互。反向残差块通常被用于轻量级的模型设计中，以提高模型的效率和推理速度。

不同模型的效果主要来源于高效算子 𝐹 的具体形式，考虑到轻量化和易用性，本文将 MMB 中的 𝐹 建模为Expanded Window MHSA（EW-MHSA）和Depth-Wise Convolution（DW-Conv）的级联，兼顾动态全局建模和静态局部信息融合的优势，同时能够有效地增加模型的感受野，提升对于下游任务的能力。

进一步，作者将 𝐹 设置为 4 并替换 DeiT 和 PVT 中标准 Transformer 结构以评估 iRMB 性能，如表3所述，可以发现iRMB可以在相同的训练设置下以更少的参数和计算获得更高的性能。

🚀二、具体添加方法

2.1 添加顺序 

（1）models/common.py    -->  加入新增的网络结构

（2）     models/yolo.py       -->  设定网络结构的传参细节，将iRMB类名加入其中。（当新的自定义模块中存在输入输出维度时，要使用qw调整输出维度）
（3） models/yolov5*.yaml  -->  新建一个文件夹，如yolov5s_iRMB.yaml，修改现有模型结构配置文件。（当引入新的层时，要修改后续的结构中的from参数）
（4）         train.py                -->  修改‘--cfg’默认参数，训练时指定模型结构配置文件 

2.2 具体添加步骤  

第①步：在common.py中添加iRMB模块 

将下面的iRMB代码复制粘贴到common.py文件的末尾。

  1. # iRMB
  2. import math
  3. from functools import partial
  4. from einops import rearrange
  5. from timm.models.layers.activations import *
  6. from timm.models.layers import DropPath
  7. from timm.models.efficientnet_builder import _parse_ksize
  8. from timm.models.efficientnet_blocks import num_groups, SqueezeExcite as SE
  10. # ========== 1.LayerNorm2d类：实现对输入张量进行二维的 Layer Normalization 操作 ==========
  11. class LayerNorm2d(nn.Module):
  13.     def __init__(self, normalized_shape, eps=1e-6, elementwise_affine=True):
  14.         super().__init__()
  15.         self.norm = nn.LayerNorm(normalized_shape, eps, elementwise_affine)
  17.     def forward(self, x):
  18.         x = rearrange(x, 'b c h w -> b h w c').contiguous()
  19.         x = self.norm(x)
  20.         x = rearrange(x, 'b h w c -> b c h w').contiguous()
  21.         return x
  24. def get_norm(norm_layer='in_1d'):
  25.     eps = 1e-6
  26.     norm_dict = {
  27.         'none': nn.Identity,
  28.         'in_1d': partial(nn.InstanceNorm1d, eps=eps),
  29.         'in_2d': partial(nn.InstanceNorm2d, eps=eps),
  30.         'in_3d': partial(nn.InstanceNorm3d, eps=eps),
  31.         'bn_1d': partial(nn.BatchNorm1d, eps=eps),
  32.         'bn_2d': partial(nn.BatchNorm2d, eps=eps),
  33.         # 'bn_2d': partial(nn.SyncBatchNorm, eps=eps),
  34.         'bn_3d': partial(nn.BatchNorm3d, eps=eps),
  35.         'gn': partial(nn.GroupNorm, eps=eps),
  36.         'ln_1d': partial(nn.LayerNorm, eps=eps),
  37.         'ln_2d': partial(LayerNorm2d, eps=eps),
  38.     }
  39.     return norm_dict[norm_layer]
  41. def get_act(act_layer='relu'):
  42.     act_dict = {
  43.         'none': nn.Identity,
  44.         'sigmoid': Sigmoid,
  45.         'swish': Swish,
  46.         'mish': Mish,
  47.         'hsigmoid': HardSigmoid,
  48.         'hswish': HardSwish,
  49.         'hmish': HardMish,
  50.         'tanh': Tanh,
  51.         'relu': nn.ReLU,
  52.         'relu6': nn.ReLU6,
  53.         'prelu': PReLU,
  54.         'gelu': GELU,
  55.         'silu': nn.SiLU
  56.     }
  57.     return act_dict[act_layer]
  59. # ========== 2.ConvNormAct类：实现卷积、规范化和激活操作的集合 ==========
  60. class ConvNormAct(nn.Module):
  62.     def __init__(self, dim_in, dim_out, kernel_size, stride=1, dilation=1, groups=1, bias=False,
  63.                  skip=False, norm_layer='bn_2d', act_layer='relu', inplace=True, drop_path_rate=0.):
  64.         super(ConvNormAct, self).__init__()
  65.         self.has_skip = skip and dim_in == dim_out
  66.         padding = math.ceil((kernel_size - stride) / 2)
  67.         self.conv = nn.Conv2d(dim_in, dim_out, kernel_size, stride, padding, dilation, groups, bias)
  68.         self.norm = get_norm(norm_layer)(dim_out)
  69.         self.act = get_act(act_layer)(inplace=inplace)
  70.         self.drop_path = DropPath(drop_path_rate) if drop_path_rate else nn.Identity()
  72.     def forward(self, x):
  73.         shortcut = x
  74.         x = self.conv(x)
  75.         x = self.norm(x)
  76.         x = self.act(x)
  77.         if self.has_skip:
  78.             x = self.drop_path(x) + shortcut
  79.         return x
  81. # ========== 3.iRMB类：反向残差注意力机制 ========== 
  82. class iRMB(nn.Module):
  84.     def __init__(self, dim_in, dim_out, norm_in=True, has_skip=True, exp_ratio=1.0, norm_layer='bn_2d',
  85.                  act_layer='relu', v_proj=True, dw_ks=3, stride=1, dilation=1, se_ratio=0.0, dim_head=64, window_size=7,
  86.                  attn_s=True, qkv_bias=False, attn_drop=0., drop=0., drop_path=0., v_group=False, attn_pre=False,inplace=True):
  87. '''
  88.    dim_in: 输入特征的维度。
  89.    dim_out: 输出特征的维度。
  90.    norm_in: 是否对输入进行标准化。
  91.    has_skip: 是否使用跳跃连接。
  92.    exp_ratio: 扩展比例。
  93.    norm_layer: 标准化层的类型。
  94.    act_layer: 激活函数的类型。
  95.    v_proj: 是否对V进行投影。
  96.    dw_ks: 深度可分离卷积的卷积核大小。
  97.    stride: 卷积的步幅。
  98.    dilation: 卷积的膨胀率。
  99.    se_ratio: SE 模块的比例。
  100.    dim_head: 注意力头的维度。
  101.    window_size: 窗口大小。
  102.    attn_s: 是否使用注意力机制。
  103.    qkv_bias: 是否在注意力机制中使用偏置。
  104.    attn_drop: 注意力机制中的dropout比例。
  105.    drop: 全连接层的dropout比例。
  106.    drop_path: DropPath 的比例。
  107.    v_group: 是否对 V 进行分组卷积。
  108.    attn_pre: 是否将注意力机制应用到输入之前。
  109.    inplace: 是否原地执行操作。
  110. '''
  112.         super().__init__()  # 调用父类的构造函数
  113.         self.norm = get_norm(norm_layer)(dim_in) if norm_in else nn.Identity()  # 条件判断，返回一个标准化层（例如 BatchNorm、LayerNorm 等）或使用空操作
  114.         dim_mid = int(dim_in * exp_ratio)  # 计算中间维度大小
  115.         self.has_skip = (dim_in == dim_out and stride == 1) and has_skip  # 条件判断，是否使用跳跃连接
  116.         self.attn_s = attn_s  # 是否使用空间注意力机制的标志
  118.         # 如果使用注意力机制
  119.         if self.attn_s:  
  120.             assert dim_in % dim_head == 0, 'dim should be divisible by num_heads'  # 确保输入维度 dim_in 可以被 dim_head 整除
  121.             self.dim_head = dim_head  # 设置每个头的维度为 dim_head
  122.             self.window_size = window_size  # 设置窗口大小
  123.             self.num_head = dim_in // dim_head  # 计算头数 self.num_head
  124.             self.scale = self.dim_head ** -0.5  # 计算缩放因子 self.scale，用于调节注意力分数
  125.             self.attn_pre = attn_pre  # 设定是否在注意力机制之前重新排列数据 self.attn_pre
  127.             # 创建 QK 卷积层、V 卷积层、注意力机制的 dropout 等
  128.             self.qk = ConvNormAct(dim_in, int(dim_in * 2), kernel_size=1, bias=qkv_bias, norm_layer='none',
  129.                                   act_layer='none')  
  130.             self.v = ConvNormAct(dim_in, dim_mid, kernel_size=1, groups=self.num_head if v_group else 1, bias=qkv_bias,
  131.                                  norm_layer='none', act_layer=act_layer, inplace=inplace)  
  132.             self.attn_drop = nn.Dropout(attn_drop)  
  134.         # 如果不使用注意力机制
  135.         else:  
  136.            # 如果需要进行 V 投影，则创建 V 卷积层；否则使用 nn.Identity() 空操作
  137.             if v_proj:  # 如果使用V投影
  138.                 self.v = ConvNormAct(dim_in, dim_mid, kernel_size=1, bias=qkv_bias, norm_layer='none',
  139.                                      act_layer=act_layer, inplace=inplace)  # 创建V卷积层
  140.             else:
  141.                 self.v = nn.Identity()  # 使用空操作
  142.         self.conv_local = ConvNormAct(dim_mid, dim_mid, kernel_size=dw_ks, stride=stride, dilation=dilation,
  143.                                       groups=dim_mid, norm_layer='bn_2d', act_layer='silu', inplace=inplace)  # 创建局部卷积层
  144.         self.se = SE(dim_mid, rd_ratio=se_ratio, act_layer=get_act(act_layer)) if se_ratio > 0.0 else nn.Identity()  # 创建空间激励模块或使用空操作
  146.         self.proj_drop = nn.Dropout(drop) 
  147.         self.proj = ConvNormAct(dim_mid, dim_out, kernel_size=1, norm_layer='none', act_layer='none', inplace=inplace)  
  148.         self.drop_path = DropPath(drop_path) if drop_path else nn.Identity()  
  150. def forward(self, x):
  151.     shortcut = x  # 保存输入的快捷连接
  152.     x = self.norm(x)  # 应用标准化层
  154.     # 提取输入 x 的形状信息
  155.     B, C, H, W = x.shape
  157.     if self.attn_s:  # 如果使用了注意力机制
  158.         # padding
  159.         if self.window_size <= 0:
  160.             window_size_W, window_size_H = W, H
  161.         else:
  162.             window_size_W, window_size_H = self.window_size, self.window_size
  163.         # 计算填充的大小
  164.         pad_l, pad_t = 0, 0
  165.         pad_r = (window_size_W - W % window_size_W) % window_size_W
  166.         pad_b = (window_size_H - H % window_size_H) % window_size_H
  167.         x = F.pad(x, (pad_l, pad_r, pad_t, pad_b, 0, 0,))  # 对输入进行填充
  168.         n1, n2 = (H + pad_b) // window_size_H, (W + pad_r) // window_size_W
  169.         x = rearrange(x, 'b c (h1 n1) (w1 n2) -> (b n1 n2) c h1 w1', n1=n1, n2=n2).contiguous()  # 重新排列输入数据
  171.         # attention
  172.         b, c, h, w = x.shape
  173.         qk = self.qk(x)  # 计算查询和键的表示
  174.         qk = rearrange(qk, 'b (qk heads dim_head) h w -> qk b heads (h w) dim_head', qk=2, heads=self.num_head,
  175.                        dim_head=self.dim_head).contiguous()  # 重排查询和键的表示
  176.         q, k = qk[0], qk[1]
  177.         attn_spa = (q @ k.transpose(-2, -1)) * self.scale  # 计算空间注意力矩阵
  178.         attn_spa = attn_spa.softmax(dim=-1)  # 对注意力矩阵进行 softmax
  179.         attn_spa = self.attn_drop(attn_spa)  # 应用注意力 dropout
  180.         if self.attn_pre:
  181.             x = rearrange(x, 'b (heads dim_head) h w -> b heads (h w) dim_head', heads=self.num_head).contiguous()  # 重排输入特征
  182.             x_spa = attn_spa @ x  # 应用注意力矩阵到输入特征
  183.             x_spa = rearrange(x_spa, 'b heads (h w) dim_head -> b (heads dim_head) h w', heads=self.num_head, h=h,
  184.                               w=w).contiguous()  # 重排输出特征
  185.             x_spa = self.v(x_spa)  # 对输出特征应用值的表示
  186.         else:
  187.             v = self.v(x)  # 计算值的表示
  188.             v = rearrange(v, 'b (heads dim_head) h w -> b heads (h w) dim_head', heads=self.num_head).contiguous()  # 重排值的表示
  189.             x_spa = attn_spa @ v  # 应用注意力矩阵到值的表示
  190.             x_spa = rearrange(x_spa, 'b heads (h w) dim_head -> b (heads dim_head) h w', heads=self.num_head, h=h,
  191.                               w=w).contiguous()  # 重排输出特征
  193.         # unpadding
  194.         x = rearrange(x_spa, '(b n1 n2) c h1 w1 -> b c (h1 n1) (w1 n2)', n1=n1, n2=n2).contiguous()  # 重新排列输出特征
  195.         if pad_r > 0 or pad_b > 0:
  196.             x = x[:, :, :H, :W].contiguous()  # 移除填充部分
  198.     else:  # 如果不使用注意力机制
  199.         x = self.v(x)  # 计算值的表示
  201.     # 应用空间激励模块和局部卷积层
  202.     x = x + self.se(self.conv_local(x)) if self.has_skip else self.se(self.conv_local(x))
  204.     # 应用输出投影的 dropout
  205.     x = self.proj_drop(x)  # 应用 dropout
  206.     x = self.proj(x)  # 应用输出投影
  208.     # 添加快捷连接并应用路径丢弃
  209.     x = (shortcut + self.drop_path(x)) if self.has_skip else x  # 添加快捷连接并应用路径丢弃
  210.     return x  # 返回处理后的结果

第②步：修改yolo.py文件 

首先找到yolo.py里面parse_model函数的这一行

加入 iRMB 这两个模块 

第③步：创建自定义的yaml文件    

  1. nc: 80  # number of classes
  2. depth_multiple: 0.33  # model depth multiple
  3. width_multiple: 0.50  # layer channel multiple
  4. anchors:
  5.   - [10,13, 16,30, 33,23]  # P3/8
  6.   - [30,61, 62,45, 59,119]  # P4/16
  7.   - [116,90, 156,198, 373,326]  # P5/32
  9. # YOLOv5 v6.0 backbone + three Attention modules
  10. backbone:
  11.   # [from, number, module, args]
  12.   [[-1, 1, Conv, [64, 6, 2, 2]],  # 0-P1/2
  13.    [-1, 1, Conv, [128, 3, 2]],    # 1-P2/4
  14.    [-1, 3, C3, [128]],
  15.    [-1, 1, Conv, [256, 3, 2]],    # 3-P3/8
  16.    [-1, 6, C3, [256]],
  17.    [-1, 1, iRMB, [256]],    # 修改1
  18.    [-1, 1, Conv, [512, 3, 2]],    # 6-P4/16
  19.    [-1, 9, C3, [512]],
  20.    [-1, 1, iRMB, [512]],    # 修改2
  21.    [-1, 1, Conv, [1024, 3, 2]],   # 9-P5/32
  22.    [-1, 3, C3, [1024]],
  23.    [-1, 1, SPPF, [1024, 5]],      # 11
  24.    [-1, 1, iRMB, [1024]],    # 修改3
  25.   ]
  27. # YOLOv5 v6.0 head
  28. head:
  29.   [[-1, 1, Conv, [512, 1, 1]],
  30.    [-1, 1, nn.Upsample, [None, 2, 'nearest']],
  31.    [[-1, 8], 1, Concat, [1]],  # cat backbone P4
  32.    [-1, 3, C3, [512, False]],  # 16
  34.    [-1, 1, Conv, [256, 1, 1]],
  35.    [-1, 1, nn.Upsample, [None, 2, 'nearest']],
  36.    [[-1, 5], 1, Concat, [1]],  # cat backbone P3
  37.    [-1, 3, C3, [256, False]],  # 20 (P3/8-small)
  39.    [-1, 1, Conv, [256, 3, 2]],
  40.    [[-1, 17], 1, Concat, [1]],  # cat head P4
  41.    [-1, 3, C3, [512, False]],  # 23 (P4/16-medium)
  43.    [-1, 1, Conv, [512, 3, 2]],
  44.    [[-1, 13], 1, Concat, [1]],  # cat head P5
  45.    [-1, 3, C3, [1024, False]],  # 26 (P5/32-large)
  47.    [[20, 23, 26], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
  48.   ]

第④步：验证是否加入成功

运行yolo.py