当前位置：首页 > news >正文

Medical-Transformer核心架构详解：Gated Axial-Attention如何革新医疗影像分析

news 2026/6/8 9:05:35

Medical-Transformer核心架构详解：Gated Axial-Attention如何革新医疗影像分析

【免费下载链接】Medical-TransformerOfficial Pytorch Code for "Medical Transformer: Gated Axial-Attention for Medical Image Segmentation" - MICCAI 2021项目地址: https://gitcode.com/gh_mirrors/me/Medical-Transformer

Medical-Transformer是一种专门针对医疗图像分割任务设计的革命性深度学习架构，它通过创新的门控轴向注意力机制，在医学影像分析领域取得了突破性进展。这个开源项目为研究人员和开发者提供了一个强大的工具，用于处理医学图像分割任务，特别是在数据量有限的情况下仍能保持出色的性能表现。

🔬 医疗图像分割的挑战与突破

传统的Transformer架构在计算机视觉任务中表现出色，但它们通常需要大规模数据集才能有效训练。然而，在医学影像领域，数据标注成本高昂且数据量有限，这成为了传统Transformer应用的主要障碍。

Medical-Transformer通过引入**Gated Axial-Attention（门控轴向注意力）**机制，成功解决了这一难题。这种创新的注意力机制不仅降低了计算复杂度，还显著提升了模型在小规模医学数据集上的表现。🎯

🏗️ Medical-Transformer整体架构

Medical-Transformer基于U-Net架构，但将传统的卷积层替换为轴向注意力模块。整个架构由编码器和解码器组成，通过跳跃连接实现多尺度特征的融合。

从上图可以看出，Medical-Transformer采用了双路径设计：全局路径处理完整图像，局部路径处理图像块。这种LoGo（Local-Global）训练策略让模型能够同时学习全局上下文信息和局部细节特征，这是其在医疗图像分割中表现出色的关键所在。

🎯 Gated Axial-Attention核心技术解析

轴向注意力机制

传统的自注意力机制在二维图像上计算时计算复杂度为O(n²)，而轴向注意力通过分别计算行和列的注意力来显著降低计算复杂度。在lib/models/axialnet.py中，轴向注意力模块的实现如下：

class AxialAttention(nn.Module): def __init__(self, in_planes, out_planes, groups=8, kernel_size=56, stride=1, bias=False, width=False): # 初始化参数 self.qkv_transform = qkv_transform(in_planes, out_planes * 2, kernel_size=1) self.bn_qkv = nn.BatchNorm1d(out_planes * 2) self.bn_similarity = nn.BatchNorm2d(groups * 3)

门控机制创新

Medical-Transformer的核心创新在于引入了门控机制，通过可学习的门控参数动态调整注意力权重。在lib/models/model_codes.py中，门控轴向注意力模块包含四个关键门控参数：

class AxialAttention_dynamic(nn.Module): def __init__(self, in_planes, out_planes, groups=8, kernel_size=56, stride=1, bias=False, width=False): # 门控参数 self.f_qr = nn.Parameter(torch.tensor(0.1), requires_grad=False) self.f_kr = nn.Parameter(torch.tensor(0.1), requires_grad=False) self.f_sve = nn.Parameter(torch.tensor(0.1), requires_grad=False) self.f_sv = nn.Parameter(torch.tensor(1.0), requires_grad=False)

这些门控参数分别控制查询-位置、键-位置、相似性-值和相似性-位置嵌入之间的交互强度，使模型能够自适应地调整不同注意力组件的贡献。

🚀 LoGo训练策略：局部与全局的完美结合

Medical-Transformer采用独特的LoGo（Local-Global）训练策略，这是其在医疗图像分割任务中取得优异性能的关键：

全局路径：处理完整的输入图像，捕获宏观结构和全局上下文信息
局部路径：将图像分割成多个小块（如32×32像素），分别处理每个小块，捕获精细的局部特征
特征融合：将全局和局部特征进行融合，获得既包含全局上下文又包含局部细节的丰富表示

在lib/models/axialnet.py的medt_net类中，这种双路径设计得到了完美实现：

# 全局路径处理 x1 = self.layer1(x) x2 = self.layer2(x1) # 局部路径处理（分块） for i in range(0,4): for j in range(0,4): x_p = xin[:,:,32*i:32*(i+1),32*j:32*(j+1)] # 处理每个图像块

📊 项目文件结构与核心模块

Medical-Transformer项目结构清晰，主要包含以下核心文件：

模型实现文件：
- lib/models/axialnet.py - 轴向注意力网络核心实现
- lib/models/model_codes.py - 各种注意力变体实现
- lib/models/utils.py - 工具函数和辅助模块
训练与测试：
- train.py - 模型训练脚本
- test.py - 模型测试脚本
- utils.py - 通用工具函数
数据加载与评估：
- lib/build_dataloader.py - 数据加载器
- metrics.py - 评估指标计算

🎨 架构优势与创新点

1. 计算效率优化 ⚡

轴向注意力机制将二维注意力分解为两个一维注意力操作，显著降低了计算复杂度，使模型能够在有限的医疗数据上有效训练。

2. 门控机制自适应 🎛️

通过门控参数动态调整不同注意力组件的权重，模型能够根据输入数据的特点自适应地调整特征提取策略。

3. 多尺度特征融合 🔄

LoGo策略实现了全局与局部特征的有机结合，使模型能够同时捕获宏观结构和微观细节。

4. 医疗数据友好 🏥

专门针对医疗图像数据量有限的特点进行优化，在小数据集上也能获得稳定可靠的性能。

🔧 快速开始指南

环境配置

项目使用Python 3.6.10和PyTorch 1.4.0，可以通过以下命令快速配置环境：

conda env create -f environment.yml conda activate medt

数据准备

按照项目要求组织数据集结构：

Train Folder/ ├── img/ │ ├── 0001.png │ └── 0002.png └── labelcol/ ├── 0001.png └── 0002.png

训练模型

使用以下命令开始训练：

python train.py --train_dataset "train_dir" --val_dataset "val_dir" \ --direc 'results' --batch_size 4 --epoch 400 \ --modelname "gatedaxialunet" --learning_rate 0.001 \ --imgsize 128 --gray "no"