济宁网站建设神华,西安网店培训班,山东网站推广营销设计,电子政务网站建设的挑战医疗图像分割任务中#xff0c;捕获多尺度信息、构建长期依赖对分割结果有非常大的影响。该论文提出了 Multi-scale Cross-axis Attention#xff08;MCA#xff09;模块#xff0c;融合了多尺度特征#xff0c;并使用Attention提取全局上下文信息。
论文地址#xff1a… 医疗图像分割任务中捕获多尺度信息、构建长期依赖对分割结果有非常大的影响。该论文提出了 Multi-scale Cross-axis AttentionMCA模块融合了多尺度特征并使用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 BaseDecodeHeaddef 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,hh,ww)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) 1self.weight nn.Parameter(torch.ones(normalized_shape))self.normalized_shape normalized_shapedef forward(self, x):sigma x.var(-1, keepdimTrue, unbiasedFalse)return x / torch.sqrt(sigma1e-5) * self.weightclass 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) 1self.weight nn.Parameter(torch.ones(normalized_shape))self.bias nn.Parameter(torch.zeros(normalized_shape))self.normalized_shape normalized_shapedef forward(self, x):mu x.mean(-1, keepdimTrue)sigma x.var(-1, keepdimTrue, unbiasedFalse)return (x - mu) / torch.sqrt(sigma1e-5) * self.weight self.biasclass 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_size1) self.conv0_1 nn.Conv2d(dim, dim, (1, 7), padding(0, 3), groupsdim)self.conv0_2 nn.Conv2d(dim, dim, (7, 1), padding(3, 0), groupsdim)self.conv1_1 nn.Conv2d(dim, dim, (1, 11), padding(0, 5), groupsdim)self.conv1_2 nn.Conv2d(dim, dim, (11, 1), padding(5, 0), groupsdim)self.conv2_1 nn.Conv2d(dim, dim, (1, 21), padding(0, 10), groupsdim)self.conv2_2 nn.Conv2d(dim, dim, (21, 1), padding(10, 0), groupsdim)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_00attn_10attn_20out2 attn_01attn_11attn_21 out1 self.project_out(out1)out2 self.project_out(out2) k1 rearrange(out1, b (head c) h w - b head h (w c), headself.num_heads)v1 rearrange(out1, b (head c) h w - b head h (w c), headself.num_heads)k2 rearrange(out2, b (head c) h w - b head w (h c), headself.num_heads)v2 rearrange(out2, b (head c) h w - b head w (h c), headself.num_heads) q2 rearrange(out1, b (head c) h w - b head w (h c), headself.num_heads) q1 rearrange(out2, b (head c) h w - b head h (w c), headself.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, headself.num_heads, hh, ww)out4 rearrange(out4, b head w (h c) - b (head c) h w, headself.num_heads, hh, ww) out self.project_out(out3) self.project_out(out4) xreturn outMODELS.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_sizeself.decoder_level Attention(in_channels[1],heads,LayerNorm_type WithBias)self.align ConvModule(in_channels[3],in_channels[0],1,conv_cfgself.conv_cfg,norm_cfgself.norm_cfg,act_cfgself.act_cfg)self.squeeze ConvModule(sum((in_channels[1],in_channels[2],in_channels[3])),in_channels[1],1,conv_cfgself.conv_cfg,norm_cfgself.norm_cfg,act_cfgself.act_cfg)self.sep_bottleneck nn.Sequential(DepthwiseSeparableConvModule(in_channels[1] in_channels[0],in_channels[3],3,padding1,norm_cfgself.norm_cfg,act_cfgself.act_cfg),DepthwiseSeparableConvModule(in_channels[3],in_channels[3],3,padding1,norm_cfgself.norm_cfg,act_cfgself.act_cfg)) def forward(self, inputs):Forward function.inputs self._transform_inputs(inputs)inputs [resize(level,sizeself.image_size,modebilinear,align_cornersself.align_corners) for level in inputs]y1 torch.cat([inputs[1],inputs[2],inputs[3]], dim1)x self.squeeze(y1) x self.decoder_level(x)x torch.cat([x,inputs[0]], dim1) x self.sep_bottleneck(x)output self.align(x) output self.cls_seg(output)return output
