咚咚

V1

2022/02/09阅读:85主题:默认主题

VT-UNet:一种用于3D医学图像肿瘤分割的Transformer模型

VT-UNet:一种用于3D医学图像肿瘤分割的Transformer模型

A Volumetric Transformer for Accurate 3D Tumor Segmentation 论文地址:https://arxiv.org/pdf/2111.13300.pdf

代码地址:https://github.com/himashi92/VT-UNet

摘要

引入主题 为体素分割设计一个计算效率高的Transformer架构是一项具有挑战性的任务。需要在编码局部和全局空间信息时保持复杂的平衡,并保存体素数据所有轴上的信息
论文方法 提议的Transformer模型具有u形编码器-解码器设计,可对输入体素进行整体处理。编码器有两个连续的自注意层来同时编码局部和全局信息,解码器基于并行移动窗口的自注意块和交叉注意块来捕获细节,并通过傅里叶位置编码来细化边界

Method

由上图可知,VT-UNet模型主要由以下几个模块组合而成:

  1. VT encoder block
  2. VT decoder block

接下来进行详细介绍

VT Encoder

3D Patch Partitioning

  1. 输入体素大小为
  2. 将其分成分成不重叠的3D patches,每个patch大小为 ,那么会得到tokens的数量为
  3. 将每个patch特征线性embedding为一个C维向量

VT Encoder Block

遵循Swin transformer的窗口设计原则,在VT编码器块中进行3D窗口操作。

提出了两种类型的窗口,即规则窗口和移位窗口,分别用VT-W-MSAVT-SW-MSA表示。

图2b为VT-W-MSA和VT-SW-MSA的设计细节,图3为窗口操作

VT-W-MSA中,将体积均匀地分割成较小的非重叠窗口,如图3所示。

由于VT-W-MSA不能使相邻窗口中的token相互看到,在VT-SW-MSA中利用了一个移位窗口(见图3中最右边的图),该窗口将VT-W-MSA相邻窗口中的token桥接起来。该窗口的灵感来自Swin Transformer

公式如下

3D Patch Merging

在每个VT Encoder Block之后,沿着空间轴以非重叠的方式合并相邻的token以产生新的token。具体来说,首先连接每组2 × 2相邻token的特征,并进行线性embedding得到通道数翻倍的特征

VT Decoder

3D Patch Expanding

  1. 首先使用线性映射将输入令牌的维数增加到原来的2倍。
  2. 经过一次reshape,将沿着空间轴重塑D/4 ×H/32 ×W/32 ×8C,从而创建D/4 ×H/16 ×W/16 ×4C ×4C。

VT Decoder Block

解码器提出了一种混合形式的SA(见图2b)。每个VT Decoder Block会收到其上一个VT Decoder Block生成的token,以及来自同一阶段的VT Encoder Block的key ( )和value ( ) tokens

VT Encoder Block有两个SA块,它们具有规则的和移位的窗口操作。VT Decoder Block具有类似的窗口操作,但使用四个SA块组成两个交叉注意(CA)模块。CA的函数可以描述为

  1. CA的右分支r作用于由之前的VT Decoder Block生成的tokens。通过其中的下标D来强调解码器的信息流。

  2. CA的左分支使用解码器生成的查询Q,以及从同一级别上VT Encoder Block获得的键K和值V。这里的思想是从编码器获取的空间信息中获益。

Fusion Module

由左右分支生成的tokens组合在一起,并馈送到下一个VT Decoder Block, 通过线性函数计算为

其中F(·)表示傅里叶特征位置编码(FPE), α控制每个CA和MSA模块的贡献。为了简单起见,选择α = 0.5

Loss Function

为了训练VT-UNet,混合最小化IDice Loss(DL)和交叉熵损失(CE)。以体素的方式修改和计算这两个损失。DL定义为

式中,H(·)和θ分别为模型和模型参数。CE定义为

因此,总分割损失为:

Experiments


推荐大家关注咚咚学AI公众号,会更新最新Cv论文和AI基本知识

分类:

人工智能

标签:

深度学习

作者介绍

咚咚
V1

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