nnU-Net-优化改进

填不完坑了,怎么这么多博客没写!!!!!


以下为口头叙述,待完善博客撰写!!!!!!

nnU-Net 虽然强,但它是纯卷积的,看东西’视野’比较窄。胰腺肿瘤太小了,需要结合整个腹部的上下文才能分得准,所以我决定引入最近很火的Mamba架构,因为它能用低显存实现全局感受野。

我没有照搬现有的U-Mamba,而是设计了一个’双向混合模块’。 这里的灵感主要有两点: 第一,医学图像不像文本,没有先后顺序,必须前后双向看,所以我把Mamba 做成了双向的。 第二,Mamba 容易忽略细节,所以我给它并联了一个卷积分支,专门抓肿瘤的边缘纹理。

效果非常超出预期。 虽然现在只跑了300 多个Epoch(总共要跑1000),但肿瘤的Dice 已经冲到了50.7%。 要知道这个数据集的SOTA(最好成绩)也就55% 左右。我们才跑了三分之一的进度就快追平SOTA 了,这说明我的’卷积+Mamba’混合策略是非常对路的。

第一部分:

  1. 关于纯 ViT 的尝试 (Task004)

最开始,我尝试过直接用 Vision Transformer 替换 nnU-Net 的 Backbone,并在 Task004 (海马体) 数据集上做了测试。

但结果并不理想,Dice 分数只有 77% 左右(通常 CNN 能达到 85%-90%)。

经过分析主要原因有两点:

第一,ViT 缺乏归纳偏置。它把图像看作序列,对局部纹理不敏感,而医学图像(尤其是海马体这种小结构)非常依赖边界的精确度。

第二,ViT 是数据饥渴型的。Task004 数据量太小,导致 Transformer 还没学好全局关系就已经过拟合了。所以纯 ViT 在这种小样本医学任务上不合适。”

  1. 关于小样本自监督的尝试

后来为了解决数据依赖问题,尝试引入小样本自监督学习 (Few-Shot Self-Supervised)。

但在实验中发现,这种方法引入了极其复杂的预训练流程,而且训练极其不稳定,很难收敛。考虑到我们要在有限算力和时间内产出稳定的改进,我认为这个方向的性价比太低,所以及时止损,暂停了这个方向。

第二部分:为什么转向Task007 与Mamba?

  1. 为什么最终选择 Task007 (胰腺)?

吸取了之前的教训,意识到需要一个既需要全局上下文,又有极高难度的任务来验证新架构。

Task004 太简单且太依赖局部,而 Task007 (胰腺及肿瘤) 是公认的数据集刷榜的:

背景干扰大: 胰腺形状极不规则,位置多变。

类别极不平衡: 胰腺肿瘤在图像里可能只占几十个像素。

这正是验证长距离建模能力的最佳场景。”

  1. 为什么选择 Bi-Hybrid Mamba?

针对 Task007 的问题,我设计了 双向混合 Mamba 架构。

这是为了解决之前 ViT 的问题:我们既要 Transformer 那样的全局视野,又要 CNN 的局部精度。

双向设计: 解决了医学图像没有时间顺序的问题,前后像素都能参考。

混合卷积: 我特意保留了一个卷积分支,专门用来抓肿瘤的边缘。

目前的实验数据非常有说服力。现在跑到了 300 多个 Epoch(总共1000):

胰腺 Dice 到了 70.1%,稳步上升。

最难的肿瘤Dice 已经冲到了50.7%。

这个方法在 007 上的核心优势在于:

相比于基线 nnU-Net,我的模型在极小病灶(肿瘤)的识别上展现出了惊人的爆发力。

肿瘤分割精度就已经逼近了 SOTA 的基线水平(约 52%-55%)。这证明了 Mamba 带来的全局上下文信息,让模型在复杂的腹部背景中,也能精准地定位到微小的肿瘤,这是传统 CNN 很难做到的。

第一,Mamba 原本是处理文本的,只能单向阅读。但 CT 图像是 3D 的,后面的像素对前面的判断也很重要。所以我设计了双向扫描,,能同时参考上下文,解决了空间信息丢失的问题。

第二,Mamba 擅长看大概轮廓,但容易把边缘磨平。而胰腺肿瘤的边缘非常模糊且不规则。所以我特意旁路并联了一个卷积层。利用卷积对纹理的高敏感性,专门用来抠肿瘤的边缘细节。

Licensed under CC BY-NC-SA 4.0