即插即用模块反向残留移动块(iRMB)详解(论文+代码)

目录

一、摘要

二、模型概要

2.1通用效率模型标准

[2.2Meta Mobile Block](#2.2Meta Mobile Block)

[2.3Micro Design: Inverted Residual Mobile Block](#2.3Micro Design: Inverted Residual Mobile Block)

[2.4Macro Design of EMO for Dense Prediction](#2.4Macro Design of EMO for Dense Prediction)

三、实验

四、iRMB代码详解


论文:https://arxiv.org/pdf/2301.01146.pdf

代码:https://github.com/zhangzjn/EMO

一、摘要

本文的重点是在权衡参数、FLOPs和性能的同时,为密集预测开发现代、高效、轻量级的模型。倒立残差块(IRB)是轻量级cnn的基础结构,但在基于注意力的研究中还没有相应的基础结构。本文从高效IRB和Transformer的有效组件的统一角度重新思考轻量级基础架构,将基于cnn的IRB扩展到基于注意力的模型,并抽象出一个用于轻量级模型设计的单残留元移动块(MMB)。根据简单而有效的设计准则,我们推导出了一种现代的反向残差移动块(Inverted Residual Mobile Block,iRMB),并构建了一个只有iRMB的类resnet高效模型(EMO)用于下游任务。在ImageNet-1K, COCO2017和ADE20K基准测试上的大量实验证明了我们的EMO优于最先进的方法,例如,EMO- 1m /2M/5M达到71.5,75.1和78.4 Top-1,超过了等阶CNN /基于注意力的模型,同时很好地权衡了参数,效率和精度:运行速度比iPhone14上的EdgeNeXt快2.8-4.0倍。

二、模型概要

2.1通用效率模型标准

在为移动应用程序设计高效的可视化模型时,我们在主观上和经验上都主张一个高效的模型应该尽可能满足以下标准:

**1.可用性:**简单的实现,不使用复杂的操作符,易于针对应用程序进行优化。

**2.一致性:**尽可能少的核心模块,以减少模型复杂性和加速部署。

**3.有效性:**具有良好的分类和密集预测性能。

**4.效率:**参数少,计算量少,精度折衷。

2.2Meta Mobile Block

以图像输入X(∈R C×H×W)为例,MMB首先使用输出/输入比为λ的扩展MLPe扩展通道维数:

然后,中间操作符F进一步增强图像特征,例如恒等操作符、静态卷积、动态MHSA等。考虑到MMB适用于高效的网络设计,我们将F作为高效操作符的概念,表示为:

最后,一个反向输入/输出比等于λ的收缩MLPs来收缩通道尺寸:

其中,使用残差连接来获得最终的输出:

2.3Micro Design: Inverted Residual Mobile Block

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

iRMB 中的 F 被建模为级联的 EW-MHSA 和DW-Conv卷积:

2.4Macro Design of EMO for Dense Prediction

基于上述标准,我们设计了一个基于一系列irmb的类似resnet的四相高效模型(EMO),用于密集应用,如图2-右所示。

1)整体框架而言,EMO仅由irmb组成,没有多样化的模块,这在设计思想上偏离了最近的高效方法。

2)对于特定模块,iRMB仅由标准卷积和多头自关注组成,没有其他复杂的操作。此外,得益于DW-Conv, iRMB可以适应跨步的下采样操作,并且不需要任何位置嵌入来引入MHSA - mapping的感应偏置。

3)对于不同的设置,我们采用逐步增加的扩展速率和通道数,详细配置如表4所示。

三、实验

这里做了相关思路整理

四、iRMB代码详解

import math  # 导入数学库
from functools import partial  # 从functools模块导入partial,用于固定函数参数值
from einops import rearrange  # 从einops库导入rearrange函数,用于重新排列张量
from timm.models.layers.activations import *  # 从timm库中导入所有激活层
from timm.models.layers import DropPath  # 从timm库中导入DropPath,用于实现随机路径丢弃
from timm.models.efficientnet_builder import _parse_ksize  # 导入解析核大小的函数
from timm.models.efficientnet_blocks import num_groups, SqueezeExcite as SE  # 导入组数计算和压缩激励类

# ========== 通用层定义区 ==========
class LayerNorm2d(nn.Module):  # 定义二维层归一化类
    def __init__(self, normalized_shape, eps=1e-6, elementwise_affine=True):  # 初始化函数
        super().__init__()  # 调用父类初始化
        self.norm = nn.LayerNorm(normalized_shape, eps, elementwise_affine)  # 创建层归一化实例

    def forward(self, x):  # 前向传播函数
        x = rearrange(x, 'b c h w -> b h w c').contiguous()  # 重排输入张量的维度,使之适应层归一化的输入需求
        x = self.norm(x)  # 应用层归一化
        x = rearrange(x, 'b h w c -> b c h w').contiguous()  # 恢复张量的原始维度排列
        return x  # 返回处理后的张量

def get_norm(norm_layer='in_1d'):  # 定义函数以获取不同类型的归一化层
    eps = 1e-6  # 设置一个小的epsilon值,防止除以零
    norm_dict = {  # 定义归一化层的字典
        # 各种归一化层和对应的构造函数,使用partial固定一些参数值
        'none': nn.Identity,  # 不进行任何操作的层
        'in_1d': partial(nn.InstanceNorm1d, eps=eps),
        'in_2d': partial(nn.InstanceNorm2d, eps=eps),
        'in_3d': partial(nn.InstanceNorm3d, eps=eps),
        'bn_1d': partial(nn.BatchNorm1d, eps=eps),
        'bn_2d': partial(nn.BatchNorm2d, eps=eps),
        'bn_3d': partial(nn.BatchNorm3d, eps=eps),
        'gn': partial(nn.GroupNorm, eps=eps),
        'ln_1d': partial(nn.LayerNorm, eps=eps),
        'ln_2d': partial(LayerNorm2d, eps=eps),
    }
    return norm_dict[norm_layer]  # 返回选定的归一化层构造函数

def get_act(act_layer='relu'):  # 定义函数以获取不同类型的激活层
    act_dict = {  # 定义激活层的字典
        # 各种激活函数及其对应的构造函数
        'none': nn.Identity,
        'sigmoid': Sigmoid,
        'swish': Swish,
        'mish': Mish,
        'hsigmoid': HardSigmoid,
        'hswish': HardSwish,
        'hmish': HardMish,
        'tanh': Tanh,
        'relu': nn.ReLU,
        'relu6': nn.ReLU6,
        'prelu': PReLU,
        'gelu': GELU,
        'silu': nn.SiLU
    }
    return act_dict[act_layer]  # 返回选定的激活层构造函数

class ConvNormAct(nn.Module):  # 定义卷积、归一化和激活组合类
    def __init__(self, dim_in, dim_out, kernel_size, stride=1, dilation=1, groups=1, bias=False,
                 skip=False, norm_layer='bn_2d', act_layer='relu', inplace=True, drop_path_rate=0.):
        super(ConvNormAct, self).__init__()  # 调用父类初始化
        self.has_skip = skip and dim_in == dim_out  # 判断是否执行跳跃连接
        padding = math.ceil((kernel_size - stride) / 2)  # 计算卷积填充
        self.conv = nn.Conv2d(dim_in, dim_out, kernel_size, stride, padding, dilation, groups, bias)  # 创建卷积层
        self.norm = get_norm(norm_layer)(dim_out)  # 创建归一化层
        self.act = get_act(act_layer)(inplace=inplace)  # 创建激活层
        self.drop_path = DropPath(drop_path_rate) if drop_path_rate else nn.Identity()  # 创建路径丢弃或空操作

    def forward(self, x):  # 前向传播函数
        shortcut = x  # 保存输入以用于跳跃连接
        x = self.conv(x)  # 应用卷积
        x = self.norm(x)  # 应用归一化
        x = self.act(x)  # 应用激活
        if self.has_skip:
            x = self.drop_path(x) + shortcut  # 应用路径丢弃并合并跳跃连接
        return x  # 返回输出

class iRMB(nn.Module):  # 定义一个继承自nn.Module的模块类iRMB

    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,
                 inplace=True):
        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')  # 创建QK卷积层,不使用归一化和激活
            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)  # 创建V卷积层,选择性分组
            self.attn_drop = nn.Dropout(attn_drop)  # 创建注意力机制的dropout
        else:  # 如果不使用注意力机制
            if v_proj:  # 如果使用V投影
                self.v = ConvNormAct(dim_in, dim_mid, kernel_size=1, bias=qkv_bias, norm_layer='none',
                                     act_layer=act_layer, inplace=inplace)  # 创建V卷积层
            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)  # 创建输出dropout
        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:  # 如果使用注意力机制
            # 计算填充尺寸
            if self.window_size <= 0:  # 如果窗口大小不大于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()  # 重排张量以适应注意力机制
            # 执行注意力操作
            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)  # 应用softmax
            attn_spa = self.attn_drop(attn_spa)  # 应用注意力dropout
            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)  # 应用V层
            else:
                v = self.v(x)  # 计算V输出
                v = rearrange(v, 'b (heads dim_head) h w -> b heads (h w) dim_head', heads=self.num_head).contiguous()  # 重排V输出
                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()  # 重排结果
            # 取消填充
            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)  # 如果不使用注意力机制,直接应用V层

        x = x + self.se(self.conv_local(x)) if self.has_skip else self.se(self.conv_local(x))  # 应用局部卷积和空间激励,根据情况添加跳跃连接

        x = self.proj_drop(x)  # 应用输出dropout
        x = self.proj(x)  # 应用输出投影

        x = (shortcut + self.drop_path(x)) if self.has_skip else x  # 如果有跳跃连接,合并结果并应用路径丢弃
        return x  # 返回输出

参考:大佬

相关推荐
bryant_meng24 分钟前
【python】OpenCV—Image Moments
开发语言·python·opencv·moments·图片矩
车载诊断技术1 小时前
电子电气架构 --- 什么是EPS?
网络·人工智能·安全·架构·汽车·需求分析
KevinRay_1 小时前
Python超能力:高级技巧让你的代码飞起来
网络·人工智能·python·lambda表达式·列表推导式·python高级技巧
跃跃欲试-迪之1 小时前
animatediff 模型网盘分享
人工智能·stable diffusion
Captain823Jack1 小时前
nlp新词发现——浅析 TF·IDF
人工智能·python·深度学习·神经网络·算法·自然语言处理
被制作时长两年半的个人练习生1 小时前
【AscendC】ReduceSum中指定workLocal大小时如何计算
人工智能·算子开发·ascendc
资源补给站2 小时前
大恒相机开发(2)—Python软触发调用采集图像
开发语言·python·数码相机
Captain823Jack2 小时前
w04_nlp大模型训练·中文分词
人工智能·python·深度学习·神经网络·算法·自然语言处理·中文分词
Black_mario2 小时前
链原生 Web3 AI 网络 Chainbase 推出 AVS 主网, 拓展 EigenLayer AVS 应用场景
网络·人工智能·web3
PieroPc2 小时前
Python 自动化 打开网站 填表登陆 例子
运维·python·自动化