咚咚

V1

2022/01/12阅读:40主题:默认主题

自监督框架USST:统一2D和3D预训练医学图像分类/分割框架

自监督框架USST:统一2D和3D预训练医学图像分类/分割框架

论文提出了一个自监督框架USST,可以同时训练2D和3D医学图像,根据输入图像维度切换到2D或3D patch embedding,并通过Transformer模型进行建模,从而使USST能够从2D和3D图像中学习表示。

论文地址:Unified 2D and 3D Pre-training for Medical Image classification and Segmentation

methods

USST的目标是使用二维和三维无标签医学图像来学习一般的图像表示,这些表示可转移和泛化到下游的二维和三维医学图像分析任务

上图为USST模型示意图。

设从混合的2D和3D数据池中进行采样来获取未标注图像 作为输入。

数据增强模块 对输入图像 进行增强操作生成两个增强视图,分别表示为 ,然后将其输入到学生-教师双路径模块中进行特征表示学习。

为了适应二维和三维医学图像,我们构建了一个 Pyramid Transformer U-Net (PTU)作为主干网络,每个block主要由SPE层和Transformer层组成。

每个SPE层根据输入图像的维度切换到特定的patch embedding

每个Transformer层由2D和3D输入共享。由于有了这些共享Transformer,USST可以从2D或3D数据中获取图像信息。

使用标准的交叉熵损失来最大化学生和教师输出之间的一致性。此外,为了最大限度地利用三维体素信息,我们还引入了slice-volume一致性损失,鼓励USST结合局部slice和全局volume信息进行训练,从而有利于从体素图像中学习更丰富的特征表示。

PTU主干网络

PTU是一个U型结构模型,每一层是由Transformer block组成,而Transformer block由SPE layer和Transformer layer组合而成,如上图所示。

SPE layer

SPE层是用来获取特定尺寸的嵌入序列(三维和二维尺寸不一致),该嵌入序列可以根据输入图像的尺寸切换到2D或3D patch嵌入。因此,SPE层在输入为2D图像时开启2D patch embedding,在输入为3D volume时开启3D patch embedding。

编码器和解码器中的SPE是不同的。编码器SPE包含一个可切换的卷积实例归一化(IN)-LeakyReLU块,它降低了特征分辨率。相比之下,解码器SPE包含一个可切换的转置卷积层,从而提高了特征分辨率

PTU Encoder

PTU编码器由一个SPE层和四个阶段的Transformer模块组成。每个Transformer模块包括一个SPE层和两个到四个Transformer层(参见上图)。

SPE:每个SPE层对输入特征进行下采样,并生成特定维数的嵌入序列。值得注意的是,在初始嵌入序列的基础上增加了一个额外的可学习的SSL token。SSL token类似于ViT中的类token,它能够通过自我注意从其他嵌入中收集知识。另外,结果序列添加可学习的位置嵌入,然后重复叠加Transformer layer

Transformer:每个Transformer层包括一个自注意层和一个具有两个隐藏层的前馈网络(FFN)。为了降低计算成本,并使PTU能够灵活地处理高分辨率特征,采用了**spatial-reduction attention(SRA)**层。给定一个查询q、一个键k和一个值v作为输入,SRA首先降低k和v的空间分辨率,然后将q、reduced k和reduced v输入到MSA层产生refine特征。

通过这种架构,PTU编码器能够在不同的尺度上学习高分辨率的特性。因此,与只能处理低分辨率特征的ViT-like net相比,它能更容易、更好地适应各种下游任务,如分类和分割

PTU Decoder

PTU解码器由三个模块块组成。在每个膜块中,输入特征图首先由一个SPE层上采样,然后由三到四个Transformer层细化。此外,还在编码器和解码器之间添加了跳过连接,以保留更多的底层细节。由于Transformer的强大能力,该解码器比纯CNN解码器具有更好的解码长远依赖性能

USST Training

所提出的USST框架是基于流行的学生-教师范式。学生路径包括PTU 和Projector ,其中 是n层多层感知器(MLP)头,θ是该路径的权值集合。

教师路径的结构与学生路径相同,但其权值μ更新为学生路径权值的展向移动平均(EMA),如下所示

在训练期间,λ使用cosine schedule从0.996增加到1。

在每个路径中,在PTU解码器的最后阶段后提取SSL token,并将其馈送给Projector,产生该路径的输出 ,其形式表示如下

使f1和f2的一致性最大化,2D和3D计算方式存在不同

Objective for 2D Images

对于二维训练图像,构建了一个标准的交叉熵损失来最小化学生和教师路径输出的误差,即 。为了避免模型坍塌,我们还对教师输出进行centering和sharpening操作(与DINO模型类似)。目标函数可以正式表示为

其中 为教师输出的中心, 为温度参数,softm(·)为sofmax函数。在教师路径中,将温度τt设置为小值,以便对其输出进行锐化。中心 首先计算当前教师epoch的数学期望,然后用EMA更新,以考虑不同批次,如下所示

其中ω是速率参数。从 中减去C可以使模型更趋均匀分布,比锐化效果好4倍。

类似地,当向学生路径输入 ,向教师路径输入 时,目标函数定义为 。因此,二维训练图像的总目标函数为

Objective for 3D Images

目前的SSL方法通常从局部slice或全局volume中学习三维医学图像的空间上下文。为了从三维训练图像中学习丰富的表示,提出以下slice-volume一致性机制。设一幅三维图像的两个增强视图分别用 表示,每个增强视图有m个二维切片。将它们分别应用于教师路径和学生路径,可以得到全局的体积度量输出

同时,还将每个增强视图的m个切片批量堆叠,作为2D输入计算切片输出,然后将所有切片输出的average处理为整体切片输出 ,如下图所示

其中 为叠加操作。在此之后,建立以下目标函数

and

其中 鼓励分别从全局volume和局部slice学习空间上下文。 鼓励增强全局volume和局部slice之间的一致性。三维训练图像的总目标函数如下所示

该函数不仅考虑了三维图像中可用的全局volume和局部slice信息,还学习了全局volume和局部slice之间的相关性。因此,它能够帮助网络从3D医学图像中捕获丰富的表示。

Optimization Scheme

引入了一个二维和三维的交替训练方案。在每个υ training step中,顺序地从2D或3D医疗图像数据集中采样几批训练数据,用于前向计算和梯度向后传播(见下图)。这个交替的训练过程将不会停止,直到USST收敛。

该训练方案具有两个优点:(1)克服了同时使用2D和3D图像的困难;(2)由于模型在每个υ step中都是基于二维或三维数据进行更新,因此可以减少二维和三维数据分布差异引起的不稳定性。

Experimental Details

预训练的设置

将输入2D patch的大小设置为224 ×224, 3D patch的大小设置为16 ×96 ×96,目的是在为SSL保留足够的信息和将计算和空间复杂性降低到一个可承受的水平之间进行权衡。

数据增强包括颜色抖动、高斯模糊/噪声、随机裁剪、缩放和翻转输入以产生两个视图。

采用了AdamW优化器和余弦衰减学习速率,10个epochs的热身期来训练USST。实验设定初始学习速率为0.0008,批量数为192,最大epoch为200,速率参数ω为0.9,温度参数τt和τs分别为0.04和0.1。

使用8个NVIDIA V100 gpu对USST进行预训练大约花了2.5

下游训练设置

分类方面,提取了预先训练好的PTU编码器,并增加了一个FC层,其输出通道作为类别数量来进行预测。

对于分割,我们在去除SSL令牌的同时,取预训练的PTU编码器和解码器,并附加一个分割头进行预测。该头包括转置卷积层、2D/3D Conv- IN-LeakyReLU和2D/3D卷积层,其核大小为1,输出通道为类别数。

分割性能由Dice系数来衡量。

分类性能是由AUC曲线下的面积来衡量的。

分类:

人工智能

标签:

图像处理

作者介绍

咚咚
V1

哈尔滨工业大学-计算机视觉