MOGANET-SA模組

iceeci發表於2024-11-10

paper
`
import torch.nn as nn
import torch
import torch.nn.functional as F
def build_act_layer(act_type):
"""Build activation layer."""
if act_type is None:
return nn.Identity()
assert act_type in ['GELU', 'ReLU', 'SiLU']
if act_type == 'SiLU':
return nn.SiLU()
elif act_type == 'ReLU':
return nn.ReLU()
else:
return nn.GELU()
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

class MultiOrderDWConv(nn.Module):
"""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.
"""

def __init__(self,
             embed_dims,
             dw_dilation=[1, 2, 3,],
             channel_split=[1, 3, 4,],
            ):
    super(MultiOrderDWConv, self).__init__()
    '''
    1/8  3/8  4/8
    1/8 dim  3/8 dim  4/8 dim
    '''
    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 = nn.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],
    )
    # DW conv 1
    self.DW_conv1 = nn.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],
    )
    # DW conv 2
    self.DW_conv2 = nn.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],
    )
    # a channel convolution
    self.PW_conv = nn.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

class MultiOrderGatedAggregation(nn.Module):
"""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_type (str): The activation type for Spatial Block.
        Defaults to 'SiLU'.
"""

def __init__(self,
             embed_dims,
             attn_dw_dilation=[1, 2, 3],
             attn_channel_split=[1, 3, 4],
             attn_act_type='SiLU',
             attn_force_fp32=False,
            ):
    super(MultiOrderGatedAggregation, self).__init__()

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

    # activation for gating and value
    self.act_value = build_act_layer(attn_act_type)
    self.act_gate = build_act_layer(attn_act_type)

    # 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

def forward_gating(self, g, v):
    with torch.autocast(device_type='cuda', enabled=False):
        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

if name == 'main':
input = torch.randn(1, 64, 32, 32).cuda()# 輸入 B C H W
block = MultiOrderGatedAggregation(embed_dims=64).cuda()
output = block(input)
print(input.size())
print(output.size())
`

相關文章