[论文阅读] TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation

论文地址：https://arxiv.org/abs/2102.04306
代码：https://github.com/Beckschen/TransUNet
发表于：arXiv 21.02

Abstract

医学图像分割是发展医疗系统的一个重要前提，尤其是在疾病诊断和治疗计划方面。在各种医学图像分割任务中，U形结构，或者说U-Net，已经成为事实上的标准，并取得了巨大的成功。然而，由于卷积操作的内在局部性，U-Net在明确建模长距离依赖方面通常表现出不足。为序列到序列预测而设计的Transformer作为替代架构出现，具有先天的全局自注意力机制，但由于低层细节不足，可能导致局部化能力有限。在本文中，我们提出了TransUNet，它同时具有Transformer和U-Net的优点，是医学图像分割的一个强有力的选择。一方面，Transformer对来自卷积神经网络(CNN)特征图的标记化图像块进行编码，作为提取全局上下文的输入序列。另一方面，解码器对编码后的特征进行升采样，然后与高分辨率的CNN特征图相结合，以实现精确定位。我们认为，Transformer可以作为医学图像分割任务的强大编码器，与U-Net相结合，通过恢复局部空间信息来增强更精细的细节。TransUNet在不同的医学应用中取得了比各种竞争方法更出色的表现，包括多器官分割和心脏分割。代码和模型可见https://github.com/Beckschen/TransUNet。

I. Motivation

本文的出现时间相对较早，但其实思路也非常直观。从一个角度看，可以将Transformer理解为一个更强的backbone，那么我们就可以用transformer去替换原有的ResNet、VGG backbone，从而来去提升网络的性能，并不需要去改动已有的CNN decoder；另一方面也可以认为，Transformer虽然全局上下文提取能力较好，但是局部细节处理能力是不如CNN的，因此依然需要保留CNN decoder架构。

从本文的标题来看，作者的思路倾向于前者。

II. Network Architecture

其实从这张图上看，可以发现甚至Encoder部分依然主要还是CNN，只是在CNN Encoder之后去接了个Transformer Encoder。也就是说，这里transformer实际上的输入是特征图而非原始图像。

文章还指出了一个有意思的细节，如果向skip connection的过程中加入transformer，性能会有进一步的提升。