CT图像通用分割模型SAMCT:任务指示提示机制与复现实践
做医学影像自动分割的人恐怕都经历过这种尴尬换一个新任务就得重新标注一批数据、重训一个专用模型有时候连同一个部位图像只是来自不同厂商的设备模型表现就直线下滑。去年我们在调研通用分割基础模型在CT图像上的适配方案时注意到一个很有意思的方向SAMCTSegment Any CT Allowing Labor-Free Task-Indicator Prompts。这个名字拆开看就是允许免劳动任务指示提示的CT任意分割。它想解决的不是分割得准不准这一个点而是能不能用一套参数尽量不动标注、少改配置就切到任意一个CT器官/结构分割任务上。我当时的第一反应是这可能才是医学影像分割工程化落地真正缺的那块拼图。这篇文章把我们从立项、方案拆解到复现踩坑的完整过程写出来希望能给做医学影像AI的同学一些参考。1. 项目起意CT分割任务的通用钥匙从哪来1.1 我们为什么需要一个CT任意分割的框架先说说背景。医学影像里的CT分割和自然图像分割差别非常大。自然图像里做个通用分割模型用海量图片训练一次然后让它框一下就能把猫抠出来这个逻辑在医学场景基本不成立。因为医学影像的数据标注贵、样本量小、类别差异大而且不同任务之间形态结构差别极大——肝脏和肺结节血管和骨骼解剖结构和病灶它们的特征分布几乎不在一个维度上。传统做法是每接一个新项目就从头训一个U-Net或者用预训练权重做迁移学习。这套流程跑多了你会发现大量的时间其实耗在数据准备、标签整理、模型调参这些重复劳动上。我们当时的想法很直接能不能仿照通用分割模型的套路做一个专门面向CT影像的通用分割器它不针对某个器官而是支持任意任务——只要给一个任务指示Task-Indicator它就自动知道现在该分割什么。这个思路一旦成立意味着同一个模型可以服务十几个科室需求每个新任务只需要少量提示甚至不需要人工精细标注。这也是SAMCT整个项目最核心的出发点。1.2 Task-Indicator Prompt这个思路是怎么冒出来的自然也避不开讨论通用分割模型比如SAM带来的范式冲击。SAM的提示机制很漂亮你给点、给框、给掩码它就能分割对应目标而且泛化性好得惊人。但在CT上直接复现问题就暴露了通用分割模型是基于自然图像的语义来理解什么是猫、什么是狗、什么是桌子而CT图像本身是灰度密度图不同组织之间边缘模糊它不具备医学先验。你让它框一下肝脏它其实不知道肝脏长什么样它只知道框里那个和周围不一样的连续区域。如果框只框了肝脏的一部分它甚至可能把肝内血管当成分割目标。所以我们想到一个关键转换任务指示不应该依赖用户手动去框选一个目标而应该是一个结构化、语义化的提示。比如一个可学习的任务Token或者一段任务描述文本告诉模型现在你的任务是分割肝脏肝脏在CT里表现为某一密度范围内的连续软组织区域。这样提示生成就变成了任务配置而不是人工标注这就为Labor-Free免劳动提供了基础。简单说这个项目的思路可以概括为三层递进第一层把CT分割从任务专用模型升级为任务条件化模型。第二层把人工提示从框、点、草图升级为任务指示Token 自动提示生成。第三层把部署形态从一个任务一个模型收敛为一个模型全部任务。这三层每一层都值得单独拆开聊。2. 核心设计解析怎么让一套模型听懂多个任务2.1 整体架构选型3D还是2.5D这是个问题开始动手之前我们首先卡在架构选择上。CT影像本质上是三维体数据最理想的情况肯定是用3D网络直接建模将所有切片视为一个整体。3D模型对空间结构关系的表达能力明显更强但它有一个致命伤显存消耗巨大而且训练数据要求高。我们当时手头有几个公开的CT分割数据集规模都不算大直接上3D网络很容易陷入过拟合再加上后期要部署到科室工作站GPU算力有限3D方案的落地成本太不可控。折中方案是2.5D。所谓2.5D就是取一个关键切片以及它的上下若干邻近切片叠成一个多通道输入。这样既保留了部分空间上下文又把模型复杂度控制在了一个合理水平。最后我们决定以2.5D作为主骨架同时保留了一个可选的3D变体用于离线批量处理。骨干网络方面我们在U-Net和Transformer之间犹豫了挺久。U-Net的好处是结构成熟、训练稳定医学影像社区里大量公开权重都基于它出了问题容易排查。Transformer的全局建模能力更强对长距离依赖的表达更好但训练不稳定小数据集上容易欠拟合。最后走了混合路线编码器使用预训练的卷积模型抽取底层特征解码器引入轻量级注意力模块做跨切片融合。这不算什么惊艳的设计但在实践里是性能与稳定性之间最好的平衡点。2.2 任务指示编码把要分割什么变成模型能懂的数字这是整个项目里最关键的模块也是我们反复迭代最多的部分。任务指示编码的目标很明确把当前任务要求分割什么结构这个信息编码成一个或多个向量注入到分割网络中让网络输出受这个条件约束。我们尝试过几种方案第一种方案是任务ID Embedding。这是最简单的做法就是给每个任务分配一个离散编号比如0代表肝脏1代表脾脏然后查表映射到向量空间。这种方案的优点是干净直接训练很快但它完全无法泛化到没见过的新任务。你训练的时候有50个任务买来的模型就只能做这50个第51个任务是没法推理的这和我们任意分割的目标相悖。第二种方案是文本描述编码。用CLIP这类视觉-语言模型把任务描述肝脏、胰腺肿瘤编码成特征向量再注入到网络中。好处是具备一定的语义泛化能力但坏处也很明显医学名词非常专业胰头钩突、肾上腺内侧支这类描述通用语言模型的理解深度是不够的而且推理时还需要依赖外部编码器流程重。第三种方案是可学习Task Token 类别语义先验的混合编码。我们为每个已知任务学习一个专属Token同时把这个任务的类别文本描述通过一个轻量的Text Encoder转成语义向量两者拼接后送入一个Prompt Injection模块。这样设计的原因在于专属Token负责记住任务之间的差异性语义向量负责提供类别的可解释提示两者互补。实测下来这种混合编码方式在已知任务上比纯Token方案平均高出一截而在新任务的零样本测试中虽然不能完全达到有监督水平但已经具备可用的雏形。最终的任务指示编码流程大致是def encode_task_indicator(task_id, task_text): # task_token: 可学习的任务专属向量 (1, d_model) task_token task_embedding_table(task_id) # text_feat: 文本编码器输出的语义向量 (1, d_model) text_feat text_encoder(task_text) # 拼接后过一层MLP得到注入网络的indicator indicator prompt_mlp(torch.cat([task_token, text_feat], dim-1)) return indicator这里有一个细节值得注意提示向量不是加到Feature Map上就完事了而是在多个解码器层级上分别注入。我们对比过只在瓶颈层注入和在不同分辨率解码层都注入的效果后者在目标比较小的时候比如小肺结节效果好得多。原因很好理解小目标的高层语义信息不充分低层细节信息在解码时又需要任务信号的引导所以提示信号必须贯穿整个解码过程。2.3 Labor-Free的提示生成机制没有人工标注提示从哪里来SAMCT里标注的Labor-Free我们认为真正含义是在推理阶段不需要用户再去手动抠细节。传统SAM你还要拿个框子去框目标。我们这套机制则希望用户只需要做一次性的任务配置之后模型自己会算出来该看哪里。那自动提示是怎么来的我们设计了一个先粗后细的两阶段流程。第一阶段模型先做一个粗分割。用窗口级的下采样特征输出一个分辨率很低比如原图的1/8的概率图这个概率图大致标出目标结构所在区域。第二步根据这个粗概率图用后处理算法提取出若干个种子点或者一个包围框作为内部提示。第二阶段把之前提取到的提示和原始图像一起送入解码器做精分割。表面看这是多了一步但好处非常大粗分割为模型提供了一个不需要人工介入的注意力先验让精分割阶段可以把信息集中在目标区域避免在全图范围内做盲目计算。这套机制的工程价值是科室医生的工作变成选择任务点开始而不是逐张图像手动画标注。我印象很深的一次临床沟通中一位医生看到我们演示时说如果每次做新项目都能这么省事那我们愿意一周多测二十个病例。这话很直白但确实是这类系统的核心吸引力。2.4 为什么选择多任务共享参数 提示条件化而不是多任务独立分支还有个方向性决策想多说两句。早期我们讨论过要不要做多任务学习的独立分支模式主干网络共享但每个任务各自接一个专用的预测头。这样做的好处是任务之间互不干扰部署时按需加载对应分支。但问题在于每个分支仍然需要各自的标注数据和训练轮次本质上没有摆脱重复训练的困境而且共享主干提取的通用特征能否适配所有任务也是个未知数。提示条件化的做法则不同。它让模型在同一个特征空间里通过任务指示向量把当前需求拧到对应的分割模式上。也就是说网络学到的不再是肝的样子或脾的样子而是在给定条件下如何从图像中提取与条件相符的区域。这样一来新增任务时我们不需要重训模型只需要准备少量提示信息甚至可以通过文本描述进行冷启动。我们做过一个对比实验同样的数据量下条件化模型在已知任务上的平均Dice比独立分支模式高1到2个百分点而在跨任务泛化上优势更明显。个别新任务的零样本Dice甚至能达到有监督训练的70%以上。3. 关键实现细节几个决定成败的小点3.1 CT特异性预处理不是所有归一化都叫适合CTCT图像的特殊性决定了预处理不能照搬自然图像。CT值本质上是人体组织对X射线衰减系数的量化单位是亨氏单位HU。不同组织的HU范围相对固定比如空气约是-1000水是0骨组织通常超过400。很多通用的预处理会直接做0到1的归一化这在自然图像上没问题放到CT上就是暴殄天物因为你会把软组织、脂肪、空气之间的密度对比关系全部打乱。我们的做法是对每个任务定义它的感兴趣窗口。比如分割肝脏时窗口通常设置为[-200, 200]HU左右超过这个范围的像素直接截断这样软组织细节能最大化保留。然后把窗口值线性映射到[0, 1]区间作为网络输入。这套方案在医学影像社区里叫窗宽窗位调整是阅片医生的基本功搬到深度学习里效果同样显著。不过这里有个坑不同任务的窗宽窗位是不同的。你把肝脏窗口设成肺结节的窗口很多细节就丢失了。所以预处理流程要和任务指示绑定。我们在输入Pipeline里维护了一张任务-窗口映射表每次迭代时根据当前任务自动选择对应的窗口参数。这看似是个小优化但实测对模型收敛速度和最终精度的影响非常直接。3.2 Prompt和视觉特征的融合位置别只在最后临门一脚很多做条件生成的朋友可能都体验过这样一个问题条件向量明明已经注入了模型好像没听见一样输出完全不受约束。我做了大量排查后确认问题基本出在注入位置上。如果只在解码器最后一层把任务指示向量拼接到特征图上模型前面的层级全部是在无提示状态下提取特征信息损失已经发生最后的拼接只能是亡羊补牢。正确做法是让提示信号在网络的前向传播中反复起作用。我们的设计是在每个解码器块里都加入一个轻量的Fusion Layer。具体来说就是def decoder_block(features, indicator, cross_slice_context): # 空间注意力根据indicator调整特征响应 gamma spatial_attn(features, indicator) features features * gamma # 与跨切片上下文信息融合 features cross_attention(features, cross_slice_context) return features这个结构有效是因为它把任务信号从一次性输入变成了持续性引导。每个解码器阶段网络都能根据任务指示有侧重地强化与目标结构相关的纹理、边界和形状特征。我们把这个设计放到整个项目里看它虽然不是最耀眼的模块但却是管用最关键的工程保障。3.3 训练策略多任务联合训练怎么让损失平衡多任务联合训练的经典问题就是任务间抢地盘。有的任务类别占比大自然在损失函数里碾压小任务有的目标边界模糊梯度噪声大会让模型整体震荡。我们试过几种常见方案直接加权、按类别频率加权以及使用任务不确定度加权的策略。直接加权简单但粗暴我们很快放弃了。按类别频率加权有一定效果但依然会让难任务拖慢训练节奏。最后采用的是基于同方差不确定性的自适应加权让网络为每个任务学习一个可调节的权重参数损失自动倾向于调节到均衡状态。公式上不打算在这里推了但直觉理解就是如果一个任务的损失波动很大说明它的不确定性高自动给它降权如果一个任务损失很稳就适当升权让模型尽快收敛。另外一个重要细节是Task Sampling策略。我们不是完全随机地在所有任务上均匀采样而是采用了难度排序采样每个epoch开头按上一个epoch中各任务的Dice升序排序对较难的任务给予更高的采样概率。这有点像课程学习先学容易的再慢慢加大难任务的比重。实测下来这个策略对整体mean Dice的提升大约有1.5个点算是一笔性价比很高的投入。3.4 推理时的自提示流程一个完整的Pipeline训练完成后推理阶段的完整流程大致是这样的输入CT序列 任务ID/任务文本 1. 加载模型和任务指示编码模块 2. 对CT序列做窗口化预处理提取关键切片及上下邻居 3. 第一阶段粗分割得到低分辨率概率图 4. 后处理概率图二值化 - 连通域分析 - 提取种子点/包围框 5. 第二阶段精分割粗提示 关键切片输入解码器 6. 后处理融合3D区域生长/多数投票 - 输出分割掩码 7. 按任务要求计算相应体积/统计指标这套流程运行在科室GPU工作站上单个病例大约在十秒级完成。和那种几十层堆叠、一次forward就要半分钟的大模型比这个速度在临床应用里是能接受的。4. 复现与实操指南从零把整个系统跑起来4.1 环境配置与依赖选择我们在复现过程中环境踩了一些坑这里直接给出稳定的组合供参考Python 3.10PyTorch 2.1CUDA 11.8及以上MONAI强烈建议用它提供的预处理和网络模块在医学影像处理上比手写轮子稳定太多einops处理维度变换时离不开albumentations或者monai.transforms二选一别混着用混着用容易出维度匹配错误显存方面8GB级别的显卡可以跑入门配置比如2.5D输入、单卡batch为8。但要做3D变体训练建议至少24GB显存否则只能堆梯度累积。4.2 数据准备从原始DICOM到训练张量CT数据大多以DICOM格式存在一个病例就是几十到几百张二维切片。我们写了一个批量转换脚本流程是用pydicom读取DICOM文件提取像素阵列和SliceLocation等空间信息将像素值转换为标准HU值注意不同设备的RescaleSlope和RescaleIntercept不同按SliceLocation排序重采样为各向同性体素常见目标是1.0mm或1.5mm间隔抽帧生成2.5D样本这个环节有个容易出错的地方方向一致性。有些数据是横断面扫描有些是冠状位或矢状位重建如果不做方向归一化模型看到的空间结构会歪严重影响训练。我们在预处理里统一将所有体数据重采样为标准的轴状位横断面再把spacing记录下来推理时输出回原始坐标系。4.3 训练参数与配置表这里列一个我们最终使用的训练配置作为参考不要当作绝对标准配置项参数值说明输入分辨率256x256在精度与显存之间取的平衡点切片数3当前切片的上下各1层骨干网络ResNet34 U-Net解码预训练权重做初始化优化器AdamW比SGD收敛更稳初始学习率3e-4线性warmup后cosine衰减Batch Size16单卡训练轮次150提前早停观察val Dice损失函数DiceLoss FocalLossDice处理类别不平衡Focal关注难例提示注入层3个解码器层每个解码块都注入数据增强旋转15度、缩放0.9-1.1、弹性形变医学图增强需保守尺度别太大训练时间大概是一个中等规模数据集几十个病例、10个任务在单张显卡上跑十多个小时。如果数据量翻倍时间也会线性上涨这是可以接受的。4.4 评估与指标选择很多人在医学分割里只盯Dice一个指标其实不够全面。我们最后采用了三个互补指标Dice系数区域重叠度衡量分割掩码和真标注的全局重合程度95% Hausdorff距离边界准确性衡量分割边界的最大偏差临床上比Dice更敏感体素精度Pixel Accuracy直观记录逐像素分类正确率在报告实验结果时这三个指标缺一不可。Dice高但Hausdorff距离很大说明边界锯齿严重这时候要检查后处理环节的平滑操作是否过度或者不足。5. 常见问题排查与避坑实录5.1 任务指示提示失灵了怎么办症状把任务A的指示向量换到任务B上模型输出看起来却几乎没变说明模型没在听提示。排查顺序是这样第一步检查提示注入层是否存在梯度消失。提示向量经过多层MLP和注意力后如果梯度回传不到最前面相当于这个分支根本没训练。我们遇到过Embedding初始化范围过大的问题把标准分布改为以0为中心、方差0.02的初始化后问题立刻缓解。第二步检查损失里提示向量的Llama指数。如果你在损失函数里只监督最终分割结果没有对提示编码做任何约束网络完全可能学成忽略提示的捷径。我们在Projection Head中间层加了一个辅助分类损失让任务指示向量在特征空间里能直接预测任务ID这相当于给提示编码了一个有力监督的信号。第三步检查任务文本的编码质量。如果你的文本编码器是通用模型对医学名词编码不准建议在训练前用训练数据的任务词表对文本编码器做一次微调校准让每个任务描述的向量分布更加稀疏且可区分。5.2 不同CT设备的域差异怎么处理CT设备厂商、扫描协议、重建算法都会影响图像质量。同一批数据上训出来的模型换到另一家医院的数据上可能掉10个点以上的Dice。这是医学影像AI里最常见的域漂移问题。我们尝试过三种方案直方图匹配、对比度归一化、以及伪影增强。直方图匹配效果很直接把参考图像的直方图分布映射到源图像上可以明显削弱设备间的密度差异。对比度归一化比全身标准化更稳定一些。伪影增强则是把一些常见的条状伪影、金属伪影在训练时随机叠加到输入图像上模型见过的伪影种类多了对未知设备的鲁棒性自然就上来了。当然最有效的还是数据层面的多中心训练。如果能拉通几家医院的数据哪怕每个中心只有少量病例联合训练后模型的域适应能力都会有质的提升。5.3 显存不足与训练稳定性问题2.5D输入虽然比3D轻量但在多任务场景下模型参数和batch都会快速膨胀。我们的做法是固定batch size但开启梯度累积等效增大batch而不爆显存使用自动混合精度AMP通常能省一半以上显存把验证时的大尺寸输入切块推理避免验证阶段OOM训练稳定性方面最棘手的是loss震荡。多任务训练下有时候前100个epoch都风平浪静第110个epoch突然loss飙升。后来定位到原因是某个任务的标注质量很差包含大量错标病例。这提醒我们多任务框架里一个脏任务会污染全局参数。后来的解决办法是给每个任务单独做一次数据质量筛查把标注置信度低的样本剔除或降权。5.4 指标原地踏步问题可能出在评测方式这种情况通常不是模型结构问题而是评测方法不对。一个常见错误是训练时图像是3个切片叠成的2.5D输入但评测时却把每个病例的所有切片独立推理再拼接切片间的空间一致性完全丢失导致边界处Dice被显著拉低。正确的评测应该是固定滑动窗口保证上下文稳定然后再做后处理拼接。另一个坑是关于窗口设置的训练时用肝脏窗[-200, 200]推理时却用了默认的[-1024, 1024]输入分布已经和训练分布不一致了性能下降是必然的。所以我们把窗口参数固化到了任务配置里推理时自动加载杜绝了人为设置错误。6. 一些个人体会与后续思路这个项目做下来我最深的体会是通用分割的难点不在于网络结构有多复杂而在于怎么把任务条件和图像特征优雅地耦合在一起。提示机制设计得好不好直接决定了系统是看起来通用还是真的通用。还有一个经验是医学影像AI的工程化落地任何环节都不能靠先训了再说。数据质量、任务配置、窗口参数、评测流程几乎每一个常规文档里不会强调的细节最后的权重都高得惊人。尤其是多任务训练的脏数据问题我们差点以为要用更大的模型去硬扛结果排查下来只是几个异常样本在捣乱。如果后续要继续扩展我会优先考虑两个方向。一是把任务指示的粒度做得更细比如结合解剖学结构树实现器官-亚区-病灶的层级指示二是探索在线任务增量学习让部署在科室的模型在使用过程中不断吸收新任务。渐进式训练会比冷启动训练稳定得多也更贴合真实工作流中标注数据逐步积累的节奏。最后分享一个小技巧如果你只想快速验证提示条件化分割这个方向在你自己数据上是否可行不必一上来就复刻完整框架。你可以先用一个现成的医学分割基线模型在你的主任务上正常训练一个版本再改造它的解码器加入任务Token然后跑两个实验对比。如果加了任务Token之后模型在新任务的少样本迁移上出现了哪怕是几个点的提升就说明这个方向在你的业务场景里值得投入。如果完全没效果那大概率是提示注入位置或监督方式的问题而不是方向本身的问题。