AAAI 2024 | 即插即用，Conv-Former注意力模块，用卷积实现Transformer效果！

import torchfrom torch import nnimport torch.nn.functional as Ffrom mmengine.model import constant_init, kaiming_init,trunc_normal_init,normal_initfrom timm.models.layers import DropPath# 论文题目：SCTNet: Single-Branch CNN with Transformer Semantic Information for Real-Time Segmentation# 中文题目: 单分支CNN结合Transformer语义信息的实时分割网络# 论文链接：https://arxiv.org/pdf/2312.17071# 官方github：https://github.com/xzz777/SCTNet# 所属机构：华中科技大学人工智能与自动化学院国家多媒体信息智能处理技术重点实验室，美团# 关键词：实时语义分割，Transformer，单分支CNN，语义信息对齐，深度学习#BN->Conv->GELU->drop->Conv2->dropclass MLP(nn.Module):    def __init__(self,                 in_channels,                 hidden_channels=None,                 out_channels=None,                 drop_rate=0.):        super(MLP,self).__init__()        hidden_channels = hidden_channels or in_channels        out_channels = out_channels or in_channels        self.norm = nn.BatchNorm2d(in_channels, eps=1e-06)        self.conv1 = nn.Conv2d(in_channels, hidden_channels, 3, 1, 1)        self.act = nn.GELU()        self.conv2 = nn.Conv2d(hidden_channels, out_channels, 3, 1, 1)        self.drop = nn.Dropout(drop_rate)        self.apply(self._init_weights)    def _init_weights(self, m):        if isinstance(m, nn.Linear):            trunc_normal_init(m.weight, std=.02)            if m.bias is not None:                constant_init(m.bias, val=0)        elif isinstance(m, (nn.SyncBatchNorm, nn.BatchNorm2d)):            constant_init(m.weight, val=1.0)            constant_init(m.bias, val=0)        elif isinstance(m, nn.Conv2d):            kaiming_init(m.weight)            if m.bias is not None:                constant_init(m.bias, val=0)    def forward(self, x):        x = self.norm(x)        x = self.conv1(x)        x = self.act(x)        x = self.drop(x)        x = self.conv2(x)        x = self.drop(x)        return x    class ConvolutionalAttention(nn.Module):    """    The ConvolutionalAttention implementation    Args:        in_channels (int, optional): The input channels.        inter_channels (int, optional): The channels of intermediate feature.        out_channels (int, optional): The output channels.        num_heads (int, optional): The num of heads in attention. Default: 8    """    def __init__(self,                 in_channels,                 out_channels,                 inter_channels,                 num_heads=8):        super(ConvolutionalAttention,self).__init__()        assert out_channels % num_heads == 0, \            "out_channels ({}) should be be a multiple of num_heads ({})".format(out_channels, num_heads)        self.in_channels = in_channels        self.out_channels = out_channels        self.inter_channels = inter_channels        self.num_heads = num_heads        self.norm = nn.BatchNorm2d(in_channels, eps=1e-06)        self.kv =nn.Parameter(torch.zeros(inter_channels, in_channels, 7, 1))        self.kv3 =nn.Parameter(torch.zeros(inter_channels, in_channels, 1, 7))        trunc_normal_init(self.kv, std=0.001)        trunc_normal_init(self.kv3, std=0.001)        self.apply(self._init_weights)    def _init_weights(self, m):        if isinstance(m, nn.Linear):            trunc_normal_init(m.weight, std=.001)            if m.bias is not None:                constant_init(m.bias, val=0.)        elif isinstance(m, (nn.SyncBatchNorm, nn.BatchNorm2d)):            constant_init(m.weight, val=1.)            constant_init(m.bias, val=.0)        elif isinstance(m, nn.Conv2d):            trunc_normal_init(m.weight, std=.001)            if m.bias is not None:                constant_init(m.bias, val=0.)    def _act_dn(self, x):        x_shape = x.shape  # n,c_inter,h,w        h, w = x_shape[2], x_shape[3]        x = x.reshape(            [x_shape[0], self.num_heads, self.inter_channels // self.num_heads, -1])   #n,c_inter,h,w -> n,heads,c_inner//heads,hw        x = F.softmax(x, dim=3)           x = x / (torch.sum(x, dim =2, keepdim=True) + 1e-06)          x = x.reshape([x_shape[0], self.inter_channels, h, w])         return x    def forward(self, x):        """        Args:            x (Tensor): The input tensor. (n,c,h,w)            cross_k (Tensor, optional): The dims is (n*144, c_in, 1, 1)            cross_v (Tensor, optional): The dims is (n*c_in, 144, 1, 1)        """        x = self.norm(x)        x1 = F.conv2d(                x,                self.kv,                bias=None,                stride=1,                padding=(3,0))          x1 = self._act_dn(x1)          x1 = F.conv2d(                x1, self.kv.transpose(1, 0), bias=None, stride=1,                padding=(3,0))          x3 = F.conv2d(                x,                self.kv3,                bias=None,                stride=1,                padding=(0,3))         x3 = self._act_dn(x3)        x3 = F.conv2d(                x3, self.kv3.transpose(1, 0), bias=None, stride=1,padding=(0,3))         x=x1+x3        return x    class CFBlock(nn.Module):    """    The CFBlock implementation based on PaddlePaddle.    Args:        in_channels (int, optional): The input channels.        out_channels (int, optional): The output channels.        num_heads (int, optional): The num of heads in attention. Default: 8        drop_rate (float, optional): The drop rate in MLP. Default:0.        drop_path_rate (float, optional): The drop path rate in CFBlock. Default: 0.2    """    def __init__(self,                 in_channels,                 out_channels,                 num_heads=8,                 drop_rate=0.,                 drop_path_rate=0.):        super(CFBlock,self).__init__()        in_channels_l = in_channels        out_channels_l = out_channels        self.attn_l = ConvolutionalAttention(            in_channels_l,            out_channels_l,            inter_channels=64,            num_heads=num_heads)        self.mlp_l = MLP(out_channels_l, drop_rate=drop_rate)        self.drop_path = DropPath(            drop_path_rate) if drop_path_rate > 0. else nn.Identity()    def _init_weights_kaiming(self, m):        if isinstance(m, nn.Linear):            trunc_normal_init(m.weight, std=.02)            if m.bias is not None:                constant_init(m.bias, val=0)        elif isinstance(m, (nn.SyncBatchNorm, nn.BatchNorm2d)):            constant_init(m.weight, val=1.0)            constant_init(m.bias, val=0)        elif isinstance(m, nn.Conv2d):            kaiming_init(m.weight)            if m.bias is not None:                constant_init(m.bias, val=0)    def forward(self, x):        x_res = x        x = x_res + self.drop_path(self.attn_l(x))        x = x + self.drop_path(self.mlp_l(x))         return xif __name__ == '__main__':    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')    input=torch.randn(1,32,256,256).to(device)    print(input.shape)    cfb = CFBlock(32,32).to(device)    output=cfb(input)    print(output.shape)