Transformer-Unet | 如何用Transformer一步一步改进 Unet？ - 极术社区

本文提出了一种用于医学图像分析的基于Transformer和UNet的神经网络，Transformer直接处理原始图像而不是提取的特征图，性能优于Trans-Unet等网络。
作者单位:北京Zoezen机器人,北航

1简介

医学图像分割在生物医学图像分析中占有重要地位，也引起了人们的广泛关注。良好的分割结果可以帮助医生进行判断，进一步改善患者体验。

在医学图像分析的众多可用方法中，UNet是最受欢迎的神经网络之一，它通过在编码器和解码器之间添加级联来保持原始特征，这使得它在工业领域仍有广泛的应用。同时，Transformer作为一种主导自然语言处理任务的模型，现已被广泛地引入到计算机视觉任务中，并在目标检测、图像分类和语义分割等任务中取得了良好的效果。因此，Transformer和UNet的结合应该比2种方法单独工作更有效。

在本文中，作者提出了Transformer-UNet，通过在原始图像中添加Transformer Block而不是在UNet中添加Feature map，并在CT-82数据集中测试本文的网络来进行胰腺分割。在实验中，形成了一个端到端的网络，并获得了比以往许多基于Unet的算法更好的分割结果。

2本文方法

首先设计一个典型的UNet作为CNN结构，使用双线性插值作为上采样方法，max-pooling作为下采样方法。为了方便实现，作者设计了一个几乎对称的网络，它可以很容易修改注意力模块和Transformer模块。然而，在T-Unet中，编码器和解码器并不直接连接，这将在本节中解释。

图1

与NLP Transformer略有不同，如图1所示。ViT将LayerNorm放在Multi-Head Attention和MLP之前，以确保输入值不会太大而无法处理。此外，ViT保留了Vaswani等人(2017)的主要设计，如Multi-Head Self-Attention和MLP层。Dosovitskiy等人(2021)进一步添加了一个可学习的数组tensor，用于在将整个序列输入到存储在T-Unet中的Transformer之前进行位置嵌入。

作者进一步修改ViT，用ELU代替GELU作为在 Transformer MLP层的激活函数，因为作者观察到ELU在实验中表现更好。与RELU和GELU相比，ELU在Transformer中使用较少，其定义为:

作者认为ELU是有用的，因为CT图像中的负值与正值同样重要。在实验中将超参数α设为1。
用上面解释的方法，用下列方程形成Transformer模型:

图2

3复现

4损失函数

5实现细节

数据集的大小对Transformer很重要。通过对CT切片而不是整个CT序列进行处理，可以扩大数据集的大小。
基于MLP的Transformer占用了大量的图形存储空间。因此，Transformer不会大量增加权重文件的大小，因此更适合于2D图像。

因此，在实验中处理CT切片，并将TUnet与现有模型Unet、Attention Unet和TransUnet进行比较。为了使模型更好地处理数据，作者将整个图像用1024进行分割，1024是数据集中所有CT切片的近似最大绝对值。

6实验

6.1结果分析

作者的主要评价方法是多个验证指标，包括mIOU值和最终预测的Dice score。CT82数据集被分离为60/22进行训练和测试。在模型中，最低分辨率为16×16，这也适用于Unet, Attention Unet和TransUnet。

为了证明结果，作者将阈值设置为0.8(即，最终预测图中值大于0.8的像素将被视为胰腺点)，在计算mIOU和像素精度值时，不仅要考虑胰腺分割的准确性，还要考虑背景的识别。

图3

图3显示了Transformer的一个主要优点，这使得模型可以使用几个Transformer层在全局和局部进行特性提取工作。

表1显示了Unet的性能和它的方差，包括TUNet。以深层Unet模型为Backbone，本文的模型能够超越UNet及其相关网络，包括目前流行的Attention Unet。

表2显示了不同模型的大小和推理时间，本文模型并没有带来特别大的参数量和推理速度。

图4

图4显示了不同网络的可视化结果，TUnet由于使用了Transformer，能够对长距离像素对进行很好的分割，因此优于以往其他基于Unet的网络。

6.2方差分析

在实验中，选择n=16作为图像patch的大小。然而，还有许多其他选项，这表明16可能不是TUnet的理想值，进一步对n=32，32进行实验。

TUnet的另一个重要特征是deep and large Unet backbone。然而，Unet和Attention Unet在浅层模型中仍然有用。由于深度模型不像浅模型那样方便，因为它们自然需要更好的硬件，如gpu，所以进一步尝试浅模型Unet Backbone。在较浅的模型中Unet中减少了1/3层CNN，并将kernel数量减少到1/4。整个模型仍然是端到端的从头到尾地训练原始模型。

表3