ICCV2023-轻量级Backbone-Rethinking Mobile Block for Efficient Attention-based Models 论文解读与代码详解

论文地址：https://arxiv.org/abs/2301.01146
源代码地址：https://github.com/zhangzjn/EMO

主要贡献

贡献一：单残差元移动块（MMB）
作者受启发与MobileNetv2的倒残差模块IRB以及Transformer中有效的 MHSA/FFN 模块，期望将他们进行结合设计一个轻量级的Backbone。并且发现这三个模块可以以一个通用的框架进行表达，称为MMB。（其实就是想将DWconv和MHSA两个拼接一下，但是讲了很长的故事）

贡献二：现代倒置残余移动块（iRMB）
基于MMB，推导出一个简单而有效的现代倒置残余移动块（iRMB）。具体来说，iRMB仅由朴素的DW-Conv和改进的EW-MHSA组成，分别用于模拟短距离/长距离依赖性。

贡献三：高效MOdel（EMO）
构建了一个类似ResNet的高效MOdel（EMO），仅使用iRMB用于下游应用。
在分类数据集ImageNet-1K和目标检测数据集COCO上面的效果如下图所示。

贡献一：单残差元移动块（MMB）

先来说说MMB，这其实是作者发现MobileNetv2的倒残差模块IRB以及Transformer中有效的 MHSA/FFN 模块可以用一个框架统一表示，并以此作为一个高效的轻量级Backbone的基础。

一个通用的表达式子是
这里的Xs就是通过1x1卷积，中间的高效模块，1x1卷积后得到的结果。作者主要改进的也就是中间的高效模块。

贡献二：现代倒置残余移动块（iRMB）

iRMB的模块结构如下

我们可以在源代码EMO-main/down-stream-tasks/mmdetection/mmdet/models/backbones/emo.py下找到iRMB的定义
解释一下定义的参数

整体结构相关：
dim_in,dim_out：输入与输出通道数

norm_in,norm_layer=‘bn_2d’:是否在输入进iRMB结构前进行norm，这里默认定义为bn_2d

exp_ratio=1.0：中间通道的扩展比例，也就是上图第一个1x1卷积输出通道数是输入卷积的exp_ratio倍，结尾的1x1卷积输出则是输入卷积的1/exp_ratio

attn_s=True，se_ratio=0.0, dim_head=64, window_size=7: 是否使用空间自注意力机制，注意力机制相关参数，SE模块的压缩比、多头注意力中每个头的维度、窗口大小。

v_proj=True: 在attn_s为False才会生效，是否对 Value 进行投影变换。

dw_ks=3, stride=1, dilation=1: 深度可分离卷积的核大小、步长、空洞率。

qkv_bias=False, attn_drop=0., drop=0., drop_path=0.: 偏置项、注意力和普通 dropout 率、路径dropout率。

v_group=False, attn_pre=False: Value投影分组与否，选择先进行注意力操作再变换 v 或者先变换 v 再做注意力操作。

整体流程我自己画了个图，如下所示

class iRMB(nn.Module):
	def __init__(self, dim_in, dim_out, norm_in=True, has_skip=True, exp_ratio=1.0, norm_layer='bn_2d',
				 act_layer='relu', v_proj=True, dw_ks=3, stride=1, dilation=1, se_ratio=0.0, dim_head=64, window_size=7,
				 attn_s=True, qkv_bias=False, attn_drop=0., drop=0., drop_path=0., v_group=False, attn_pre=False):
		super().__init__()
		self.norm = get_norm(norm_layer)(dim_in) if norm_in else nn.Identity()
		dim_mid = int(dim_in * exp_ratio)
		self.has_skip = (dim_in == dim_out and stride == 1) and has_skip
		self.attn_s = attn_s
		if self.attn_s:
			assert dim_in % dim_head == 0, 'dim should be divisible by num_heads'
			self.dim_head = dim_head
			self.window_size = window_size
			self.num_head = dim_in // dim_head
			self.scale = self.dim_head ** -0.5
			self.attn_pre = attn_pre
			self.qk = ConvNormAct(dim_in, int(dim_in * 2), kernel_size=1, bias=qkv_bias, norm_layer='none',
								  act_layer='none')
			self.v = ConvNormAct(dim_in, dim_mid, kernel_size=1, groups=self.num_head if v_group else 1, bias=qkv_bias,
								 norm_layer='none', act_layer=act_layer, inplace=inplace)
			self.attn_drop = nn.Dropout(attn_drop)
		else:
			if v_proj:
				self.v = ConvNormAct(dim_in, dim_mid, kernel_size=1, bias=qkv_bias, norm_layer='none',
									 act_layer=act_layer, inplace=inplace)
			else:
				self.v = nn.Identity()
		self.conv_local = ConvNormAct(dim_mid, dim_mid, kernel_size=dw_ks, stride=stride, dilation=dilation,
									  groups=dim_mid, norm_layer='bn_2d', act_layer='silu', inplace=inplace)
		self.se = SE(dim_mid, rd_ratio=se_ratio, act_layer=get_act(act_layer)) if se_ratio > 0.0 else nn.Identity()

		self.proj_drop = nn.Dropout(drop)
		self.proj = ConvNormAct(dim_mid, dim_out, kernel_size=1, norm_layer='none', act_layer='none', inplace=inplace)
		self.drop_path = DropPath(drop_path) if drop_path else nn.Identity()

	def forward(self, x):
		shortcut = x
		x = self.norm(x)
		B, C, H, W = x.shape
		if self.attn_s:
			# padding
			if self.window_size <= 0:
				window_size_W, window_size_H = W, H
			else:
				window_size_W, window_size_H = self.window_size, self.window_size
			pad_l, pad_t = 0, 0
			pad_r = (window_size_W - W % window_size_W) % window_size_W
			pad_b = (window_size_H - H % window_size_H) % window_size_H
			x = F.pad(x, (pad_l, pad_r, pad_t, pad_b, 0, 0,))
			n1, n2 = (H + pad_b) // window_size_H, (W + pad_r) // window_size_W
			x = rearrange(x, 'b c (h1 n1) (w1 n2) -> (b n1 n2) c h1 w1', n1=n1, n2=n2).contiguous()
			# attention
			b, c, h, w = x.shape
			qk = self.qk(x)
			qk = rearrange(qk, 'b (qk heads dim_head) h w -> qk b heads (h w) dim_head', qk=2, heads=self.num_head,
						   dim_head=self.dim_head).contiguous()
			q, k = qk[0], qk[1]
			attn_spa = (q @ k.transpose(-2, -1)) * self.scale
			attn_spa = attn_spa.softmax(dim=-1)
			attn_spa = self.attn_drop(attn_spa)
			if self.attn_pre:
				x = rearrange(x, 'b (heads dim_head) h w -> b heads (h w) dim_head', heads=self.num_head).contiguous()
				x_spa = attn_spa @ x
				x_spa = rearrange(x_spa, 'b heads (h w) dim_head -> b (heads dim_head) h w', heads=self.num_head, h=h,
								  w=w).contiguous()
				x_spa = self.v(x_spa)
			else:
				v = self.v(x)
				v = rearrange(v, 'b (heads dim_head) h w -> b heads (h w) dim_head', heads=self.num_head).contiguous()
				x_spa = attn_spa @ v
				x_spa = rearrange(x_spa, 'b heads (h w) dim_head -> b (heads dim_head) h w', heads=self.num_head, h=h,
								  w=w).contiguous()
			# unpadding
			x = rearrange(x_spa, '(b n1 n2) c h1 w1 -> b c (h1 n1) (w1 n2)', n1=n1, n2=n2).contiguous()
			if pad_r > 0 or pad_b > 0:
				x = x[:, :, :H, :W].contiguous()
		else:
			x = self.v(x)

		x = x + self.se(self.conv_local(x)) if self.has_skip else self.se(self.conv_local(x))

		x = self.proj_drop(x)
		x = self.proj(x)

		x = (shortcut + self.drop_path(x)) if self.has_skip else x
		return x

贡献三：高效MOdel（EMO）

作者在原文中提到，一个高效模型应尽可能满足以下标准：

➀ 简单的实现，不使用复杂的运算符，并且易于针对应用程序进行优化。
➁ 均匀性。尽可能少的核心模块，以降低模型复杂性并加快部署。
➂有效性。具有良好的分类和密集预测性能。
➃ 效率更少的参数和计算与准确性的权衡。

EMO比较简单，以EMO5M为例，包含5个stage，其中stage0多一个最开始的MSPatchEmb，其它4个stage都是iRMB模块的堆叠。

其中MSPatchEmb的代码如下，就是定义最开始的几层分组空洞卷积，其中dilations=[1, 2, 3]定义了卷积的个数，c_group=-1定义了分组卷积的组数（-1为动态匹配）

class MSPatchEmb(nn.Module):

	def __init__(self, dim_in, emb_dim, kernel_size=2, c_group=-1, stride=1, dilations=[1, 2, 3],
				 norm_layer='bn_2d', act_layer='silu'):
		super().__init__()
		self.dilation_num = len(dilations)
		assert dim_in % c_group == 0
		c_group = math.gcd(dim_in, emb_dim) if c_group == -1 else c_group
		self.convs = nn.ModuleList()
		for i in range(len(dilations)):
			padding = math.ceil(((kernel_size - 1) * dilations[i] + 1 - stride) / 2)
			self.convs.append(nn.Sequential(
				nn.Conv2d(dim_in, emb_dim, kernel_size, stride, padding, dilations[i], groups=c_group),
				get_norm(norm_layer)(emb_dim),
				get_act(act_layer)(emb_dim)))

	def forward(self, x):
		if self.dilation_num == 1:
			x = self.convs[0](x)
		else:
			x = torch.cat([self.convs[i](x).unsqueeze(dim=-1) for i in range(self.dilation_num)], dim=-1)
			x = reduce(x, 'b c h w n -> b c h w', 'mean').contiguous()
		return x