【论文笔记】MCANet: Medical Image Segmentation withMulti-Scale Cross-Axis Attention

        医疗图像分割任务中，捕获多尺度信息、构建长期依赖对分割结果有非常大的影响。该论文提出了 Multi-scale Cross-axis Attention（MCA）模块，融合了多尺度特征，并使用Attention提取全局上下文信息。

论文地址：MCANet: Medical Image Segmentation with Multi-Scale Cross-Axis Attention

代码地址：https://github.com/haoshao-nku/medical_seg

一、MCA(Multi-scale Cross-axis Attention)

MCA的结构如下，将E2/3/4通过concat连接起来（concat前先插值到同样分辨率），经过1x1的卷积后（压缩通道数来降低计算量），得到了包含多尺度信息的特征图F，然后在X和Y方向使用不同大小的卷积核进行卷积运算（比如1x11的卷积是x方向，11x1的是y方向，这里可以对着代码看，容易理解），将Q在X和Y方向交换后（这就是Cross-Axis），经过注意力模块后，将多个特征图相加，并融合E1，经过卷积后得到输出。该模块有以下特点：

1、注意力机制作用在多个不同尺度的特征图；

2、Multi-Scale x-Axis Convolution和Multi-Scale y-Axis Convolution分别关注不同轴的特征，在计算注意力时交叉计算，使得不同方向的特征都能被关注到。

MCA细节如下图，输入特征图进入x和y方向的路径，经过不同大小的卷积后进行融合，然后跨轴（x和y轴的Q交换）计算Attention，最后得到输出特征图。

二、代码

MCA的代码如下所示，总体来说比较简单：

from audioop import bias
from pip import main
import torch
import torch.nn as nn
import torch.nn.functional as F
import numbers
from mmseg.registry import MODELS
from einops import rearrange
from ..utils import resize
from mmcv.cnn import ConvModule, DepthwiseSeparableConvModule
from mmseg.models.decode_heads.decode_head import BaseDecodeHead


def to_3d(x):
    return rearrange(x, 'b c h w -> b (h w) c')

def to_4d(x,h,w):
    return rearrange(x, 'b (h w) c -> b c h w',h=h,w=w)

class BiasFree_LayerNorm(nn.Module):
    def __init__(self, normalized_shape):
        super(BiasFree_LayerNorm, self).__init__()
        if isinstance(normalized_shape, numbers.Integral):
            normalized_shape = (normalized_shape,)
        normalized_shape = torch.Size(normalized_shape)

        assert len(normalized_shape) == 1

        self.weight = nn.Parameter(torch.ones(normalized_shape))
        self.normalized_shape = normalized_shape

    def forward(self, x):
        sigma = x.var(-1, keepdim=True, unbiased=False)
        return x / torch.sqrt(sigma+1e-5) * self.weight

class WithBias_LayerNorm(nn.Module):
    def __init__(self, normalized_shape):
        super(WithBias_LayerNorm, self).__init__()
        if isinstance(normalized_shape, numbers.Integral):
            normalized_shape = (normalized_shape,)
        normalized_shape = torch.Size(normalized_shape)

        assert len(normalized_shape) == 1

        self.weight = nn.Parameter(torch.ones(normalized_shape))
        self.bias = nn.Parameter(torch.zeros(normalized_shape))
        self.normalized_shape = normalized_shape

    def forward(self, x):
        mu = x.mean(-1, keepdim=True)
        sigma = x.var(-1, keepdim=True, unbiased=False)
        return (x - mu) / torch.sqrt(sigma+1e-5) * self.weight + self.bias

class LayerNorm(nn.Module):
    def __init__(self, dim, LayerNorm_type):
        super(LayerNorm, self).__init__()
        if LayerNorm_type =='BiasFree':
            self.body = BiasFree_LayerNorm(dim)
        else:
            self.body = WithBias_LayerNorm(dim)

    def forward(self, x):
        h, w = x.shape[-2:]
        return to_4d(self.body(to_3d(x)), h, w)
    
class Attention(nn.Module):
    def __init__(self, dim, num_heads,LayerNorm_type,):
        super(Attention, self).__init__()
        self.num_heads = num_heads   
        self.temperature = nn.Parameter(torch.ones(num_heads, 1, 1))   

        self.norm1 = LayerNorm(dim, LayerNorm_type)
        self.project_out = nn.Conv2d(dim, dim, kernel_size=1)      
        self.conv0_1 = nn.Conv2d(dim, dim, (1, 7), padding=(0, 3), groups=dim)
        self.conv0_2 = nn.Conv2d(dim, dim, (7, 1), padding=(3, 0), groups=dim)
        self.conv1_1 = nn.Conv2d(dim, dim, (1, 11), padding=(0, 5), groups=dim)
        self.conv1_2 = nn.Conv2d(dim, dim, (11, 1), padding=(5, 0), groups=dim)
        self.conv2_1 = nn.Conv2d(
            dim, dim, (1, 21), padding=(0, 10), groups=dim)
        self.conv2_2 = nn.Conv2d(
            dim, dim, (21, 1), padding=(10, 0), groups=dim)

    def forward(self, x):
        b,c,h,w = x.shape   
        x1 = self.norm1(x)
        attn_00 = self.conv0_1(x1)
        attn_01= self.conv0_2(x1)  
        attn_10 = self.conv1_1(x1)
        attn_11 = self.conv1_2(x1)
        attn_20 = self.conv2_1(x1)
        attn_21 = self.conv2_2(x1)   
        out1 = attn_00+attn_10+attn_20
        out2 = attn_01+attn_11+attn_21   
        out1 = self.project_out(out1)
        out2 = self.project_out(out2)  
        k1 = rearrange(out1, 'b (head c) h w -> b head h (w c)', head=self.num_heads)
        v1 = rearrange(out1, 'b (head c) h w -> b head h (w c)', head=self.num_heads)
        k2 = rearrange(out2, 'b (head c) h w -> b head w (h c)', head=self.num_heads)
        v2 = rearrange(out2, 'b (head c) h w -> b head w (h c)', head=self.num_heads)   
        q2 = rearrange(out1, 'b (head c) h w -> b head w (h c)', head=self.num_heads) 
        q1 = rearrange(out2, 'b (head c) h w -> b head h (w c)', head=self.num_heads)       
        q1 = torch.nn.functional.normalize(q1, dim=-1)
        q2 = torch.nn.functional.normalize(q2, dim=-1)
        k1 = torch.nn.functional.normalize(k1, dim=-1)
        k2 = torch.nn.functional.normalize(k2, dim=-1)          
        attn1 = (q1 @ k1.transpose(-2, -1))
        attn1 = attn1.softmax(dim=-1)   
        out3 = (attn1 @ v1) + q1      
        attn2 = (q2 @ k2.transpose(-2, -1))
        attn2 = attn2.softmax(dim=-1)   
        out4 = (attn2 @ v2) + q2                         
        out3 = rearrange(out3, 'b head h (w c) -> b (head c) h w', head=self.num_heads, h=h, w=w)
        out4 = rearrange(out4, 'b head w (h c) -> b (head c) h w', head=self.num_heads, h=h, w=w)       
        out =  self.project_out(out3)  + self.project_out(out4) + x
          
        return out

@MODELS.register_module()
class MCAHead(BaseDecodeHead):
    def __init__(self,in_channels,image_size,heads,c1_channels,
                 **kwargs):
        super(MCAHead, self).__init__(in_channels,input_transform = 'multiple_select',**kwargs)
        self.image_size = image_size
        self.decoder_level = Attention(in_channels[1],heads,LayerNorm_type = 'WithBias')
        self.align = ConvModule(
            in_channels[3],
            in_channels[0],
            1,
            conv_cfg=self.conv_cfg,
            norm_cfg=self.norm_cfg,
            act_cfg=self.act_cfg)
        self.squeeze = ConvModule(
            sum((in_channels[1],in_channels[2],in_channels[3])),
            in_channels[1],
            1,
            conv_cfg=self.conv_cfg,
            norm_cfg=self.norm_cfg,
            act_cfg=self.act_cfg)
        self.sep_bottleneck = nn.Sequential(
            DepthwiseSeparableConvModule(
                in_channels[1] + in_channels[0],
                in_channels[3],
                3,
                padding=1,
                norm_cfg=self.norm_cfg,
                act_cfg=self.act_cfg),
            DepthwiseSeparableConvModule(
                in_channels[3],
                in_channels[3],
                3,
                padding=1,
                norm_cfg=self.norm_cfg,
                act_cfg=self.act_cfg))             
    def forward(self, inputs):
        """Forward function."""
        inputs = self._transform_inputs(inputs)
        inputs = [resize(
                level,
                size=self.image_size,
                mode='bilinear',
                align_corners=self.align_corners
            ) for level in inputs]
        y1 = torch.cat([inputs[1],inputs[2],inputs[3]], dim=1)
        x = self.squeeze(y1)  
        x = self.decoder_level(x)
        x = torch.cat([x,inputs[0]], dim=1) 
        x = self.sep_bottleneck(x)
        
        output = self.align(x)  
        output = self.cls_seg(output)
        return output