IVP 2021 UNETR: Transformers for 3D Medical Image Segmentation
动机
- 具有跟踪和扩展路径的全卷积神经网络（FCNNs）（如编码器和解码器），近年来在各种医学图像分割应用中显示出了突出的地位。在这些体系结构中，编码器通过学习全局上下文表示来扮演一个不可或缺的角色，该全局上下文表示将进一步用于解码器的语义输出预测。
- 卷积层作为FCNN的主要构造块，其局部性限制了此类网络学习长时空间依赖的能力。
- 在NLP领域，基于transformer的模型已经在各种任务中实现了最先进的基准。transformer中的自我注意机制使他们能够动态地突出词序列的重要特征并学习其长时依赖关系。这一概念最近被扩展到计算机视觉，引入了视觉transformer(ViT)。在ViT中，一张图像被表示为将用于直接预测用于图像分类的类标签的图片块嵌入序列。
方法
如上图所示，UNETR采用收缩-膨胀模式，由一堆transformer作为编码器，通过跳过连接与解码器相连。
编码器。创建一个三维输入（h×W×D×C）的一维序列，通过将其划分为flatten的相同大小的非重叠图片块。使用一个线性层将平面化的图片块投影到一个K维嵌入空间中，该空间在整个transformer中保持不变。此外，为了保留所提取的图片块的空间信息，在投影图片块上加入一个一维可学习的位置嵌入。在嵌入层之后，利用由多头自注意力机制(MSA)和多层感知器(MLP)子层组成的transformer块堆叠，MLP包括两个具有GELU激活函数的线性层，MSA块包括n个平行的自注意力(SA)头。SA块是一个参数化函数，它学习输入序列(z)中的两个元素及其query和key表征之间的相似性。
编码器的多个分辨率的特征与解码器合并，从transformer中提取序列特征，将其装换为统一大小的张量。利用连续的3×3×3的卷积层，然后进行批处理归一化，将大小重塑的张量从嵌入空间投射到输入空间。在编码器的bottleneck（即transformer的最后一层输出），应用一个去卷积层到变换后的特征图，以2倍的增加其分辨率。将调整大小的特征映射与前一个transformer输出的特征映射连接起来，将它们传递到连续的3×3×3卷积层中，并使用反卷积层对输出进行上采样。对于直到原始输入分辨率的所有其他后续层，重复此过程。最终输出被传递到带有softmax激活函数的1×1×1卷积层，以生成逐像素的语义预测。损失函数是dice和交叉熵项的组合。
实验
为了涵盖不同的目标和图像模式，采用MSD挑战中任务1（脑肿瘤MRI分割）和任务9（脾脏CT分割）的数据集进行实验，并将数据拆分进行5倍交叉验证。实验证实了UNETR在CT和MR图像模式下MSD数据集的脑肿瘤和脾脏体积分割任务上的有效性，并且持续表现出良好的性能。
贡献
- UNETR是为三维分割量身定制的，并且直接利用三维数据。
- UNETR将transformer作为分割网络的主编码器，并通过跳过连接直接连接到解码器，而不是将transformer作为分割网络中的注意力层。
- UNETR不依赖于主干CNN来生成输入序列，而是直接利用标记过的图片块。