模块拼接|即插即用模块|MCANet|注意力机制

 MCANet（Multi-scale Cross-Axis Attention Network）是一个基于多尺度交叉轴注意力机制的网络模型，医疗图像分割任务中，捕获多尺度信息、构建长期依赖对分割结果有非常大的影响。 Multi-scale Cross-axis Attention（MCA）模块融合了多尺度特征，并使用Attention提取全局上下文信息。

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

代码地址：GitHub - haoshao-nku/medical_seg

以下是整合的核心模块，可以直接与其他模型缝合。

import torch
from torch import nn
import math


__all__ = ['MCALayer', 'MCAGate']

# 计算输入特征图的标准差，将标准差作为输出
class StdPool(nn.Module) :
    def __init__(self):
        super(StdPool, self).__init__()

    def forward (self, x):
        b, c, _, _ = x.size()
        # b, c, _, _ = x.size()：这行代码获取输入张量 x 的尺寸，并分别将批次大小（batch size）、通道数（channels）、高度（height）和宽度（width）赋值给 b、c、_ 和 _
        # _ 是 Python 中的占位符，表示这些值在后续代码中不会被使用
        std = x.view(b, c, -1).std(dim=2, keepdim=True)
        std = std.reshape(b, c, 1, 1)  # 使输出张量的形状与输入张量 x 的空间维度（高度和宽度）相匹配

        return std
        # 用于提取特征图中的统计信息，而不是传统的空间降采样或特征选择



# 接受不同的池化类型和卷积核大小作为输入参数，并通过组合这些池化操作的结果以及一个可学习的权重来生成一个门控信号。这个门控信号随后用于调整输入特征图 x
class MCAGate (nn.Module):
    # 用于动态的调整特征图 接受不同的池化类型和卷积核大小作为输入参数
    def __init__(self, k_size, pool_types=['avg', 'std']):

        super(MCAGate, self).__init__()

        self.pools = nn.ModuleList([])
        for pool_type in pool_types:
            if pool_type == 'avg':
                self.pools.append(nn.AdaptiveAvgPool2d(1))
            elif pool_type == 'max':
                self.pools.append(nn.AdaptiveMaxPool2d(1))
            elif pool_type == 'std':
                self.pools.append(StdPool())
            else:
                raise NotImplementedError

        self.conv = nn.Conv2d(1, 1, kernel_size=(1, k_size),  stride=1, padding=(0, (k_size - 1) // 2), bias=False)
        self.sigmoid = nn.Sigmoid()

        self.weight = nn.Parameter(torch.rand(2))

    def forward(self, x):
        feats = [pool(x) for pool in self.pools]

        if len(feats) == 1:
            out = feats[0]
        elif len(feats) == 2:
            weight = torch.sigmoid(self.weight)
            out = 1/2*(feats[0] + feats[1]) + weight[0] * feats[0] + weight[1] * feats[1]
        else:
            assert False, "特征提取异常"

        out = out.permute(0, 3, 2, 1).contiguous()
        out = self.conv(out)
        out = out.permute(0, 3, 2, 1).contiguous()

        out = self.sigmoid(out)
        out = out.expand_as(x)

        return x * out



class MCALayer(nn.Module):
    def __init__(self, channel, no_spatial=False):  #channel:输入特征图的通道数，no_spatial布尔值，指示是否构建空间维度的交互。
        """Constructs a MCA module.
        Args:
            inp: Number of channels of the input feature maps
            no_spatial: whether to build channel dimension interactions
        """
        super(MCALayer, self).__init__()

        lambd = 1.5
        gamma = 1
        temp = round(abs((math.log2(channel) - gamma) / lambd))
        kernel = temp if temp % 2 else temp - 1

        self.h_cw = MCAGate(3)
        self.w_hc = MCAGate(3)
        self.no_spatial = no_spatial
        if not no_spatial:
            self.c_hw = MCAGate(kernel)

    def forward(self, x):
        x_h = x.permute(0, 2, 1, 3).contiguous()
        x_h = self.h_cw(x_h)
        x_h = x_h.permute(0, 2, 1, 3).contiguous()  # 将输入x在高度和宽度维度上进行转置，然后通过h_cw门控机制进行变换，最后再转置回原始维度。

        x_w = x.permute(0, 3, 2, 1).contiguous()
        x_w = self.w_hc(x_w)
        x_w = x_w.permute(0, 3, 2, 1).contiguous()  # 将输入x在宽度和高度维度上进行转置，通过w_hc门控机制进行变换，然后再转置回原始维度。

        if not self.no_spatial:   # 如果no_spatial为False，则对原始输入x应用c_hw门控机制，这可能在空间维度上执行某种形式的卷积或注意力操作。
            x_c = self.c_hw(x)
            x_out = 1 / 3 * (x_c + x_h + x_w)
        else:
            x_out = 1 / 2 * (x_h + x_w)

        return x_out