【Moga Net】Efficient Multi-order Gated Aggregation Network 阅读

Efficient Multi-order Gated Aggregation Network

目录

  • Efficient Multi-order Gated Aggregation Network
      • 导读
    • 背景
    • 网络结构
    • 消融实验
    • 特征可视化
    • 实验结果
    • 官方代码库

导读

今天分享一篇发在arxiv上的backbone基础网络创新论文【MogaNet】,作者来自于西湖大学和浙江大学。MogaNet 将DNN的表征能力表示为lens of interaction complexities。MogaNet为纯卷积网络,通过类似于人类视觉系统的多阶特征聚集模块同时具备局部感知和全局上下文聚集能力,通过特殊设计的SMixer和CMixer 分别从Spatial维度和Channel维度聚集上下文信息,实现了计算复杂度-性能间的平衡,在分类、检测、分割任务上取得了不错性能指标。
【Moga Net】Efficient Multi-order Gated Aggregation Network 阅读_第1张图片
【Moga Net】Efficient Multi-order Gated Aggregation Network 阅读_第2张图片

背景

卷积层利用区域密集连接和平移不变性约束对输入图像的邻域相关性进行编码。卷积网络通过级联的卷积层逐渐增加感受野,自适应的识别语义特征。但是卷积网络学习到的特征更关注于图像局部纹理特征,缺乏对图像全局上下文信息的获取。通过宏观网络设计(ASPP,FPN)和增加全局上下文aggregation (SE-Block,GC-Block) 可以部分解决该问题。

基于自注意力的Transformer更关注于图像全局依赖,在很多领域都超过了CNN的性能,
但是Transformer的二次方计算复杂度限制了网络的计算效率和在下游密集预测任务上的应用。由于缺少CNN的归纳偏置特性,Transformer学习图像相关邻域关系能力较弱。最近有研究将CNN中的FPN金字塔结构和重新引入ViT。

人类视觉系统不仅提取局部特征,同时也聚集全局感知特征,比深度学习更紧凑和高效。大多数现代DNN倾向于编码极低或极高复杂性的相互作用,而不是信息最丰富的中间作用。
【Moga Net】Efficient Multi-order Gated Aggregation Network 阅读_第3张图片
如图所示,人类视觉特征在遮挡50%左右时即可获取几乎全部图像的信息,而DNN在10%以下/90%以上时获取更多的信息。

网络结构

MogaNet 提出的宏观网络结构。经过stem调整输入维度,SMixer获取空间特征,CMixer 混合channel。
【Moga Net】Efficient Multi-order Gated Aggregation Network 阅读_第4张图片

【Moga Net】Efficient Multi-order Gated Aggregation Network 阅读_第5张图片
【Moga Net】Efficient Multi-order Gated Aggregation Network 阅读_第6张图片

整体网络结构

消融实验

【Moga Net】Efficient Multi-order Gated Aggregation Network 阅读_第7张图片
【Moga Net】Efficient Multi-order Gated Aggregation Network 阅读_第8张图片
【Moga Net】Efficient Multi-order Gated Aggregation Network 阅读_第9张图片
【Moga Net】Efficient Multi-order Gated Aggregation Network 阅读_第10张图片
【Moga Net】Efficient Multi-order Gated Aggregation Network 阅读_第11张图片

【Moga Net】Efficient Multi-order Gated Aggregation Network 阅读_第12张图片
【Moga Net】Efficient Multi-order Gated Aggregation Network 阅读_第13张图片

特征可视化

实验结果

【Moga Net】Efficient Multi-order Gated Aggregation Network 阅读_第14张图片

官方代码库

https://github.com/Westlake-AI/openmixup

MogaNet

ElementScale

class ElementScale(nn.Module):
    """A learnable element-wise scaler."""

    def __init__(self, embed_dims, init_value=0., requires_grad=True):
        super(ElementScale, self).__init__()
        self.scale = nn.Parameter(
            init_value * torch.ones((1, embed_dims, 1, 1)),
            requires_grad=requires_grad
        )

    def forward(self, x):
        return x * self.scale

Moga Block

class MogaBlock(BaseModule):
    """A block of MogaNet.
    Args:
        embed_dims (int): Number of input channels.
        ffn_ratio (float): The expansion ratio of feedforward network hidden
            layer channels. Defaults to 4.
        drop_rate (float): Dropout rate after embedding. Defaults to 0.
        drop_path_rate (float): Stochastic depth rate. Defaults to 0.1.
        act_cfg (dict, optional): The activation config for projections and FFNs.
            Default: dict(type='GELU').
        norm_cfg (dict): Config dict for normalization layer.
            Defaults to ``dict(type='BN')``.
        init_value (float): Init value for Layer Scale. Defaults to 1e-5.
        attn_dw_dilation (list): Dilations of three DWConv layers.
        attn_channel_split (list): The raletive ratio of splited channels.
        attn_act_cfg (str): The activation config for the gating branch.
            Default: dict(type='SiLU').
        init_cfg (obj:`mmcv.ConfigDict`): The Config for initialization.
            Default: None.
    """

    def __init__(self,
                 embed_dims,
                 ffn_ratio=4.,
                 drop_rate=0.,
                 drop_path_rate=0.,
                 act_cfg=dict(type='GELU'),
                 norm_cfg=dict(type='BN', eps=1e-5),
                 init_value=1e-5,
                 attn_dw_dilation=[1, 2, 3],
                 attn_channel_split=[1, 3, 4],
                 attn_act_cfg=dict(type='SiLU'),
                 attn_force_fp32=False,
                 init_cfg=None):
        super(MogaBlock, self).__init__(init_cfg=init_cfg)
        self.out_channels = embed_dims

        self.norm1 = build_norm_layer(norm_cfg, embed_dims)[1]

        # spatial attention
        self.attn = MultiOrderGatedAggregation(
            embed_dims,
            attn_dw_dilation=attn_dw_dilation,
            attn_channel_split=attn_channel_split,
            attn_act_cfg=attn_act_cfg,
            attn_force_fp32=attn_force_fp32,
        )
        self.drop_path = DropPath(
            drop_path_rate) if drop_path_rate > 0. else nn.Identity()

        self.norm2 = build_norm_layer(norm_cfg, embed_dims)[1]

        # channel MLP
        mlp_hidden_dim = int(embed_dims * ffn_ratio)
        self.mlp = ChannelAggregationFFN(  # DWConv + Channel Aggregation FFN
            embed_dims=embed_dims,
            feedforward_channels=mlp_hidden_dim,
            act_cfg=act_cfg,
            ffn_drop=drop_rate,
        )

        # init layer scale
        self.layer_scale_1 = nn.Parameter(
            init_value * torch.ones((1, embed_dims, 1, 1)), requires_grad=True)
        self.layer_scale_2 = nn.Parameter(
            init_value * torch.ones((1, embed_dims, 1, 1)), requires_grad=True)

    def forward(self, x):
        # spatial
        identity = x
        x = self.layer_scale_1 * self.attn(self.norm1(x))
        x = identity + self.drop_path(x)
        # channel
        identity = x
        x = self.layer_scale_2 * self.mlp(self.norm2(x))
        x = identity + self.drop_path(x)
        return x

SMixer

class MultiOrderGatedAggregation(BaseModule):
   """Spatial Block with Multi-order Gated Aggregation.
   Args:
       embed_dims (int): Number of input channels.
       attn_dw_dilation (list): Dilations of three DWConv layers.
       attn_channel_split (list): The raletive ratio of splited channels.
       attn_act_cfg (dict, optional): The activation config for FFNs.
           Default: dict(type='GELU').
       init_cfg (obj:`mmcv.ConfigDict`): The Config for initialization.
           Default: None.
   """

   def __init__(self,
                embed_dims,
                attn_dw_dilation=[1, 2, 3],
                attn_channel_split=[1, 3, 4],
                attn_act_cfg=dict(type="SiLU"),
                attn_force_fp32=False,
                init_cfg=None):
       super(MultiOrderGatedAggregation, self).__init__(init_cfg=init_cfg)

       self.embed_dims = embed_dims
       self.attn_force_fp32 = attn_force_fp32
       self.proj_1 = Conv2d(
           in_channels=embed_dims, out_channels=embed_dims, kernel_size=1)
       self.gate = Conv2d(
           in_channels=embed_dims, out_channels=embed_dims, kernel_size=1)
       self.value = MultiOrderDWConv(
           embed_dims,
           dw_dilation=attn_dw_dilation,
           channel_split=attn_channel_split,
       )
       self.proj_2 = Conv2d(
           in_channels=embed_dims, out_channels=embed_dims, kernel_size=1)

       # activation for gating and value
       self.act_value = custom_build_activation_layer(attn_act_cfg)
       self.act_gate = custom_build_activation_layer(attn_act_cfg)

       # decompose
       self.sigma = ElementScale(
           embed_dims, init_value=1e-5, requires_grad=True)

   def feat_decompose(self, x):
       x = self.proj_1(x)
       # x_d: [B, C, H, W] -> [B, C, 1, 1]
       x_d = F.adaptive_avg_pool2d(x, output_size=1)
       x = x + self.sigma(x - x_d)
       x = self.act_value(x)
       return x

   @force_fp32()
   def forward_gating(self, g, v):
       g = g.to(torch.float32)
       v = v.to(torch.float32)
       return self.proj_2(self.act_gate(g) * self.act_gate(v))

   def forward(self, x):
       shortcut = x.clone()
       # proj 1x1
       x = self.feat_decompose(x)
       # gating and value branch
       g = self.gate(x)
       v = self.value(x)
       # aggregation
       if not self.attn_force_fp32:
           x = self.proj_2(self.act_gate(g) * self.act_gate(v))
       else:
           x = self.forward_gating(self.act_gate(g), self.act_gate(v))
       x = x + shortcut
       return x

multi-orderDWConv

class MultiOrderDWConv(BaseModule):
    """Multi-order Features with Dilated DWConv Kernel.
    Args:
        embed_dims (int): Number of input channels.
        dw_dilation (list): Dilations of three DWConv layers.
        channel_split (list): The raletive ratio of three splited channels.
        init_cfg (obj:`mmcv.ConfigDict`): The Config for initialization.
            Default: None.
    """

    def __init__(self,
                 embed_dims,
                 dw_dilation=[1, 2, 3,],
                 channel_split=[1, 3, 4,],
                 init_cfg=None):
        super(MultiOrderDWConv, self).__init__(init_cfg=init_cfg)

        self.split_ratio = [i / sum(channel_split) for i in channel_split]
        self.embed_dims_1 = int(self.split_ratio[1] * embed_dims)
        self.embed_dims_2 = int(self.split_ratio[2] * embed_dims)
        self.embed_dims_0 = embed_dims - self.embed_dims_1 - self.embed_dims_2
        self.embed_dims = embed_dims
        assert len(dw_dilation) == len(channel_split) == 3
        assert 1 <= min(dw_dilation) and max(dw_dilation) <= 3
        assert embed_dims % sum(channel_split) == 0

        # basic DW conv
        self.DW_conv0 = Conv2d(
            in_channels=self.embed_dims,
            out_channels=self.embed_dims,
            kernel_size=5,
            padding=(1 + 4 * dw_dilation[0]) // 2,
            groups=self.embed_dims,
            stride=1, dilation=dw_dilation[0], #1
        )
        # DW conv 1
        self.DW_conv1 = Conv2d(
            in_channels=self.embed_dims_1,
            out_channels=self.embed_dims_1,
            kernel_size=5,
            padding=(1 + 4 * dw_dilation[1]) // 2,
            groups=self.embed_dims_1,
            stride=1, dilation=dw_dilation[1],# 1
        )
        # DW conv 2
        self.DW_conv2 = Conv2d(
            in_channels=self.embed_dims_2,
            out_channels=self.embed_dims_2,
            kernel_size=7,
            padding=(1 + 6 * dw_dilation[2]) // 2,
            groups=self.embed_dims_2,
            stride=1, dilation=dw_dilation[2],#2
        )
        # a channel convolution
        self.PW_conv = Conv2d(  # point-wise convolution
            in_channels=embed_dims,
            out_channels=embed_dims,
            kernel_size=1)

    def forward(self, x):
        x_0 = self.DW_conv0(x)
        x_1 = self.DW_conv1(
            x_0[:, self.embed_dims_0: self.embed_dims_0+self.embed_dims_1, ...])
        x_2 = self.DW_conv2(
            x_0[:, self.embed_dims-self.embed_dims_2:, ...])
        x = torch.cat([
            x_0[:, :self.embed_dims_0, ...], x_1, x_2], dim=1)
        x = self.PW_conv(x)
        return x

Channel Aggregation

class ChannelAggregationFFN(BaseModule):
    """An implementation of FFN with Channel Aggregation.
    Args:
        embed_dims (int): The feature dimension. Same as
            `MultiheadAttention`.
        feedforward_channels (int): The hidden dimension of FFNs.
        kernel_size (int): The depth-wise conv kernel size as the
            depth-wise convolution. Defaults to 3.
        ffn_drop (float, optional): Probability of an element to be
            zeroed in FFN. Default 0.0.
        init_cfg (obj:`mmcv.ConfigDict`): The Config for initialization.
            Default: None.
    """

    def __init__(self,
                 embed_dims,
                 feedforward_channels,
                 kernel_size=3,
                 act_cfg=dict(type='GELU'),
                 ffn_drop=0.,
                 init_cfg=None):
        super(ChannelAggregationFFN, self).__init__(init_cfg=init_cfg)

        self.embed_dims = embed_dims
        self.feedforward_channels = feedforward_channels
        self.act_cfg = act_cfg

        self.fc1 = Conv2d(
            in_channels=embed_dims,
            out_channels=self.feedforward_channels,
            kernel_size=1)
        self.dwconv = Conv2d(
            in_channels=self.feedforward_channels,
            out_channels=self.feedforward_channels,
            kernel_size=kernel_size,
            stride=1,
            padding=kernel_size // 2,
            bias=True,
            groups=self.feedforward_channels)
        self.act = build_activation_layer(act_cfg)
        self.fc2 = Conv2d(
            in_channels=feedforward_channels,
            out_channels=embed_dims,
            kernel_size=1)
        self.drop = nn.Dropout(ffn_drop)

        self.decompose = Conv2d(
            in_channels=self.feedforward_channels,  # C -> 1
            out_channels=1, kernel_size=1,
        )
        self.sigma = ElementScale(
            self.feedforward_channels, init_value=1e-5, requires_grad=True)
        self.decompose_act = build_activation_layer(act_cfg)

    def feat_decompose(self, x):
        # x_d: [B, C, H, W] -> [B, 1, H, W]
        x = x + self.sigma(x - self.decompose_act(self.decompose(x)))
        return x

    def forward(self, x):
        # proj 1
        x = self.fc1(x)
        x = self.dwconv(x)
        x = self.act(x)
        x = self.drop(x)
        # proj 2
        x = self.feat_decompose(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x
        ```

你可能感兴趣的:(基础网络,paper,reading,深度学习,计算机视觉,人工智能)