A Simple-yet-Effective Volume Contrastive Learning Framework for 3D Medical Image Analysis
一个简单而有效的三维医学图像分析体积对比学习框架
是通过自监督学习搭建的学习框架
自监督学习是无监督学习下的一个分支;
它通过从未标注的数据中生成伪标签来进行训练。自监督学习利用数据的内在结构或属性来创建训练任务,从而在没有人工标注的情况下进行学习。自监督学习常用于预训练模型,然后将预训练的模型迁移到有标注的数据上进行微调(finetune)。
自监督学习的监督信息不是人工标注的,而是是通过辅助任务(pretext)在大规模无监督数据中自动构造监督信息,通过得到的标签,就可以类似有监督学习一样进行训练。
VoCo的辅助任务(pretext):首先从不同区域的一组基本作物,同时强制它们之间的特征差异,我们将它们用作不同区域的类分配。然后,我们随机裁剪子体积,并通过对比它们与不同基本作物的相似性来预测它们属于哪个类(位于哪个区域),这可以看作是预测不同子体积的上下文位置。
具体流程:
1.首先定义一个Pretext task (辅助任务),即从无监督的数据中,通过巧妙地设计自动构造出有监督(伪标签)数据,学习得到一个预训练模型。
通过10k dataset 训练得到 预训练模型 VoCo_10k.pt.
2.然后可以将预训练模型,通过简单的Finetune,应用到下游的多个应用场景,以分割任务为例:BTCV(腹部)、MM-WHS(全心)。
-
加载预训练模型:
- 使用预训练模型的权重初始化下游任务的模型(SwinUNETR)。
- 在
main.py
中,通过args.pretrained_checkpoint
加载预训练模型VoCo_10k.pt
。将预训练模型VoCo_10k.pt中的预训练权重,加载到SwinUNETR模型中.
-
定义下游任务的模型结构:
- 使用SwinUNETR模型进行分割任务。加载预训练权重后,定义模型的结构。
- 在
main.py
中,定义了SwinUNETR模型,并加载预训练权重。
-
设置优化器和损失函数:
- 定义优化器(如AdamW)和损失函数(如DiceCELoss)。
- 在
main.py
中,设置了优化器和损失函数。
-
进行训练和验证:
- 使用训练数据进行模型训练,并在验证数据上进行评估。
- 在
trainer.py
中,定义了训练和验证的具体过程。
对比学习
分为原型级(prototype-level contrastive)和实例级(instance-level contrastive)
- prototype-level contrastive
- 该方法的关键步骤是生成一组原型,这些原型可以被视为数据集中不同类别的中心点或代表性特征。生成原型的途径有两种:一是通过在线聚类(online clustering),即在整个数据集上进行实时聚类,将实例分配到不同的簇中,并将每个簇的中心作为该簇的原型;二是随机初始化一组原型,然后通过训练过程中的反向传播(backpropagation)在线更新这些原型,使其逐渐接近各自代表类别的中心。
- 在获得原型之后,对比学习的目标转变为将每个输入实例与这些原型进行对比。模型被训练以区分输入实例与对应类别原型之间的相似性或差异性,而非直接与其他实例进行一对一的对比。这种以原型为中心的对比强调了不同类别间的结构关系和内在联系。
- 在训练过程中,给定一个输入实例,模型会计算其特征向量与所有原型之间的距离(如使用余弦相似度、欧氏距离等)。如果输入实例与某一原型的距离较小,表明它们在特征空间中接近,即该实例可能属于该原型代表的类别。反之,如果与某原型距离较大,则认为它们属于不同的类别。通过最大化正例对(输入实例与相应类别原型)之间的相似度,同时最小化负例对(输入实例与无关类别原型)之间的相似度,模型学习到能够在特征空间中有效分离不同类别的表示。
- VoCo(Volume Contrast)框架遵循了原型级对比学习的基本理念,但避免了耗时的在线聚类和原型更新过程。取而代之,VoCo利用3D医学图像中固有的上下文位置信息生成基础作物(base crops)作为原型(bases)。这些基础作物代表了不同位置的特征,为模型提供了一种无需额外标注即可学习到3D医学图像中高级语义(如器官位置、解剖结构等)的方法。在VoCo中,通过对比随机裁剪的子体积与这些基础作物的相似性,模型预测子体积所处的上下文位置,从而实现对上下文位置先验的隐式编码。
- instance-level contrastive
- 每个输入实例(如图像或3D医学体积)会经历一系列数据增强技术,如随机裁剪、旋转、色彩抖动或翻转。这些变换创造出同一实例的多个“视图”或“增强版本”,保持其核心内容不变但稍微改变其外观。
- 对于每个原始实例,一对增强视图被视为“正样本对”,即尽管它们表面上有所不同,但实际上代表同一个底层实例。相反,由不同实例组成的对被视为“负样本对”,表示不同的实例。
- 使用深度神经网络(通常为图像处理中的卷积神经网络或体数据专用架构)从原始和增强实例中提取特征表示。这些表示通常通过附加的投影头投射到更低维度的嵌入空间。
- 模型通过使用对比损失函数(如InfoNCE损失或其变种)进行训练,该函数促使正样本对的特征嵌入彼此靠近(即具有高相似度),同时使负样本对的嵌入相互远离。这种损失函数在学习到的表示中强加了实例特异性概念,确保即使面对视觉上相似的变体,模型也能区分独特的实例。
对于不同部位分隔的调整:
pre - processing (scale intensity,spacing) can be different.
Now we set scale intensity[-175, 250] for abdomen and head neck [0,1700] for heart [-500,1000] for chest. for spacing and size it should be depended by specific datadet
比较vit和swin-transformer
Vision Transformer (ViT)
- 基本概念:ViT 是将 Transformer 架构直接应用于图像的一种尝试。它首先将图像分割成一系列固定大小的补丁(patches),然后将这些补丁线性投影到一个嵌入空间中,最后将它们作为序列输入到标准的 Transformer 模型中进行处理。
- 补丁处理:ViT 将图像分割成大小相同的补丁,并将每个补丁视为 Transformer 输入序列中的一个元素。这种处理方式忽略了图像中的局部结构信息。
- 位置编码:ViT 需要为输入序列的每个元素(即图像补丁)添加位置编码,以保留补丁在原始图像中的空间位置信息。
- 全局自注意力:ViT 在处理图像补丁时采用全局自注意力机制,这意味着模型在处理任何补丁时都会考虑到图像中的所有补丁。这使得 ViT 对于大尺寸图像而言计算成本较高。
Swin Transformer
- 基本概念:Swin Transformer 是一种层次化的 Transformer,它通过引入了一个可变形的窗口机制来处理图像。Swin Transformer 在局部窗口内应用自注意力,然后逐渐合并这些窗口以捕获更大范围的依赖关系。
- 窗口自注意力:Swin Transformer 的关键创新之一是在局部窗口内应用自注意力,这减少了计算量并允许模型捕获局部特征。随着层次的加深,窗口会合并,使得模型能够捕获更大范围的上下文信息。
- 移位窗口机制:为了增强模型的表示能力并避免局部窗口内自注意力的限制,Swin Transformer 引入了移位窗口机制,通过在不同层之间交错窗口的位置来增强跨窗口连接。
- 层次化特征:Swin Transformer 生成了一系列层次化的特征表示,这对于多尺度或金字塔式的视觉任务特别有用。
比较总结
- 局部性 vs. 全局性:ViT 采用全局自注意力处理图像的所有部分,而 Swin Transformer 通过窗口自注意力和层次化结构强调局部性和全局性之间的平衡。
- 计算效率:Swin Transformer 通过限制自注意力的范围和层次化处理提高了计算效率,特别是对于大尺寸图像,而 ViT 的全局自注意力机制在大尺寸图像上计算成本较高。
- 适用性:Swin Transformer 的层次化和可变形的窗口机制使其在多种尺寸和类型的视觉任务上表现更为灵活和有效,而 ViT 作为一种更直接的应用 Transformer 架构于图像的尝试,可能需要更大的数据集和更多的计算资源来实现最佳性能
腹部器官(BTCV)
"labels": {
"0": "background", # 背景
"1": "spleen", # 脾脏
"2": "rkid", # 右肾
"3": "lkid", # 左肾
"4": "gall", # 胆囊
"5": "eso", # 食管
"6": "liver", # 肝脏
"7": "sto", # 胃
"8": "aorta", # 主动脉
"9": "IVC", # 下腔静脉
"10": "veins", # 静脉
"11": "pancreas", # 胰腺
"12": "rad", # 右肾上腺
"13": "lad" # 左肾上腺
},