UTNet 用于医学图像分割的混合Transformer

点击上方“计算机视觉工坊”，选择“星标”

干货第一时间送达

作者丨李响

来源丨 GiantPandaCV

文章目录

• 1 概述

• 2 浅析 Transformer 架构

  • 2.1 重看 Self-attention Mechanism

  • 2.2 高效的 Self-attention Mechanism

  • 2.3 Relative Positional Encoding

• 3 UTNet 的整体结构

• 4 实验

• 5 总结

• 6 参考链接

  1 概述

在这里插入图片描述

很久之前读了这篇接收在 MICCAI 2021 上的文章，复现调试了代码之后还没有及时整理一篇阅读笔记。由于在 MICCAI 上，这篇文章同样没有大量的实验对比工作，但是提出的方法思路清晰易懂，值得借鉴。arXiv: https://arxiv.org/abs/2107.00781 。作为前置依赖，本篇阅读笔记首先介绍了 Transformer Architecture 和在医学图像分割上的应用；其次，分析了论文中提出的 UTNet 架构（主干 U-Net，混合 Transformer 等 module）；最后，看了下论文的实验结果。

在语义分割上，FCN 这类卷积的编码器-解码器架构衍生出的模型在过去几年取得了实质性进展，但这类模型存在两个局限。第一，卷积仅能从邻域像素收集信息，缺乏提取明确全局依赖性特征的能力；第二，卷积核的大小和形状往往是固定的，因此它们不能灵活适应输入的图像或其他内容。相反，Transformer architecture 由于自注意力机制具有捕获全局依赖特征的能力，且允许网络根据输入内容动态收集相关特征。

值得注意的是，Transformer 架构的训练需要的更大的开销，因为自注意力机制（self-attention）在序列长度方面具有的时间和空间复杂度。基于此，标准的 self-attention 大多数以 patch-wise 方式应用到模型中，比如使用 16 × 16 这种小扁平图像块作为输入序列，或者在来自 CNN 主干的特征图之上对图像进行编码，这些特征图一般是下采样后的低分辨率图像。这里问题就出来了，对于医学图像分割任务目标位置敏感的特殊性，一些欠分割或者过分割的区域都在目标周围，往往需要高分辨率特征。此外，有些实验论证，在 ImageNet 上进行了预训练，Transformer 要比 ResNet 差，Transformer 真的能适应医学图像这种小数据集分割任务吗？

为了解决上面的问题，文章中提出的 U-Net 混合 Transformer 网络：UTNet，它整合了卷积和自注意力策略用于医学图像分割任务。应用卷积层来提取局部强度特征，以避免对 Transformer 进行大规模的预训练，同时使用自注意力来捕获全局特征。为了提高分割质量，还提出了一种 efficient self-attention，在时间和空间上将整体复杂度从 O(n2) 显着降低到接近 O(n)。此外，在 self-attention 模块中使用相对位置编码来学习医学图像中的内容-位置关系。

2 浅析 Transformer 架构

上面我们对 Transformer 在医学图像分割上的应用和局限做了概述，这里我们简单分析下 Transformer 架构，对这部分很了解的读者可以跳过。

2.1 重看 Self-attention Mechanism

Transformer 建立在多头自注意机制 (MHSA) 模块上，MHSA 是由多个 Self-Attention 组成的。下图是 Self-Attention 的结构，在计算的时候需要用到矩阵 **Q(查询)，K(键值)，V(值)**。在实际中，Self-Attention 接收的是输入(单词的表示向量 x 组成的矩阵 X) 或者上一个 Encoder block 的输出。而 Q,K,V 是通过 Self-Attention 的输入进行线性变换得到的。

得到矩阵 Q, K, V 之后就可以计算出 Self-Attention 的输出了，计算的公式如下。其中 d 是 Q,K 矩阵的列数（向量维度），公式中计算矩阵 Q 和 K 每一行向量的内积，为了防止内积过大，因此除以 d 的平方根。将 Q, K, V 展平并转置为大小为 n × d 的序列，其中 n = HW。P ∈ Rn×n 被命名为上下文聚合矩阵，用作权重以收集上下文信息。

通过这种方式，self-attention 本质上具有全局感受野，擅长捕捉全局依赖。此外，上下文聚合矩阵可以适应输入内容，以实现更好的特征聚合。关于更细节的内容，这里就不多介绍了。需要关注的是，n×d 矩阵的点乘会导致 O(n2d) 复杂度。通常，当特征图的分辨率很大时，n 远大于 d，因此序列长度 n 在自注意力计算中占主导地位，这使得高分辨率特征图中应用自注意力是不可行的，例如对于 16 × 16 特征图，n = 256，对于 128 × 128 特征图，n = 16384。这一点在本篇笔记的概述中提到过。

2.2 高效的 Self-attention Mechanism

那么如何降低 O(n2) 的复杂度呢？由于医学图像是高度结构化的数据，除了边界区域外，局部像素的高分辨率特征和图中的其他像素特征存在相似性，因此，所有像素之间的成对注意力计算往往是低效和冗余的。所以需要一种高效的自我注意机制，计算如下。

观察与 2.1 中公式不同之处，主要思想很简单，是将 K 和 V 矩阵做了一个 Low-dimension Embedding 达到减小计算量的目的，对应的上下文聚合矩阵 P 的 size 也会被修改。通过这样做，可以把复杂度降低到 O(nkd)。相当于，这里我们可以做任意下采样操作，例如平均/最大池化或带步长卷积操作。比如，使用 1×1 卷积，然后使用双线性插值对特征图进行下采样，缩小后的 size 是可以推断的。

基于 MHSA 的 Encoder 和 Decoder block 分别如下图所示，关于这两部分如何插入到 U-Net 中，在后面会有解释。

2.3 Relative Positional Encoding

Transformer 的相对位置编码大概作用是让像素间保持空间位置关系（从ViT的实验看，去掉 position embedding 后，性能会下降3个点以上），对于图像就是保持二维信息，它对于图像任务来讲很重要。像目前的相对位置编码设计，都算是将位置编码整合到了 Attention 的定义中去了，没有显式地使用绝对位置编码。读者们应该可以发现 2.2 的 MHSA 图里，在 softmax 之前的 pair-wise attention logit 计算中使用了像素 i 和 j 的位置编码，具体的计算如下。

相应的，self-attention 的计算公式又需要做修正了，和之前的区别就是相对宽度和高度是在 low-dimensional projection （低维投影）之后计算的。对 2.2 的高效 self-attention 修正后的计算方式如下。

3 UTNet 的整体结构

上一部分我们对 UTNet 中的 Transformer 架构做了清晰的介绍，这一部分会整体解释下 UTNet 结构，也就是如何把 Transformer 的 encoder 和 decoder 合理加入到主干 U-Net 中。

如上图所示 UTNet 结构图，整体上还是保持 U 型。(b) 是一个经典的残差块，传统的 U-Net 改进方法也是这么做的，这样也可以提高分割任务的准确率，避免网络深度带来的梯度爆炸和梯度消失等问题，这些都是老生常谈了，我们不重点关注。(c) 是一个标准的 Transformer Decoder 设计。可以发现，遵循了 U-Net 的标准设计，但将每个构建块的最后一个卷积（最高的除外）替换为 2.2 的 Transformer 模块。此外，低三层的跨层连接也被替换为了 Transformer Decoder，我认为还是很好理解的哈，关于这样做的目的在笔记前面陆陆续续都说了一些，下面总结下。

这种混合架构可以利用卷积图像的归纳偏差来避免大规模预训练，以及 Transformer 捕获全局特征关系的能力。由于错误分割的区域通常位于感兴趣区域的边界，高分辨率的上下文信息可以在分割中发挥至关重要的作用。因此，重点放在了自我注意模块上，这使得有效处理大尺寸特征图成为可能。没有将自注意力模块简单地集成到来自 CNN 主干的特征图之上，而是将 Transformer 模块应用于编码器和解码器的每个级别，以从多个尺度收集长期依赖关系。请注意，没有在原始分辨率上应用 Transformer，因为在网络的非常浅层中添加 Transformer 模块对实验没有帮助，但会引入额外的计算。一个可能的原因是网络的浅层更多地关注详细的纹理，其中收集全局上下文特征效果肯定不理想。

4 实验

数据集选用 Multi-Centre, Multi-Vendor & Multi-Disease Cardiac Image Segmentation Challenge。

下图是消融实验的结果。下表除了参数量和预测时间比 U-Net 等大，其他均是 SOTA。(a) 为不同自注意力位置的影响消融；(b) 为不同高效 Self-attention 的 projection 效果（对应 2.2）；（c）为 Transformer 编码器、Transformer 解码器和相对位置编码的影响消融。

更多的实验对比和结果可视化如下所示，大多数指标比 CBAM 要优秀。

5 总结

这篇阅读笔记大多为个人理解，代码复现我后面也许还会更新一篇文章，由于一些医学图像处理任务数据集的特殊性，vit 在医学图像上的应用还需要不断优化，最近有一些不错的想法，也欢迎交流和纠正！

6 参考链接

• https://arxiv.org/abs/2107.00781

• Campello, V.M., Palomares, J.F.R., Guala, A., Marakas, M., Friedrich, M., Lekadir, K.: Multi-Centre, Multi-Vendor & Multi-Disease Cardiac Image Segmentation Challenge (Mar 2020)

本文仅做学术分享，如有侵权，请联系删文。

干货下载与学习

后台回复：巴塞罗那自治大学课件，即可下载国外大学沉淀数年3D Vison精品课件

后台回复：计算机视觉书籍，即可下载3D视觉领域经典书籍pdf

后台回复：3D视觉课程，即可学习3D视觉领域精品课程

计算机视觉工坊精品课程官网：3dcver.com

1.面向自动驾驶领域的多传感器数据融合技术

2.面向自动驾驶领域的3D点云目标检测全栈学习路线！(单模态+多模态/数据+代码)
3.彻底搞透视觉三维重建：原理剖析、代码讲解、及优化改进
4.国内首个面向工业级实战的点云处理课程
5.激光-视觉-IMU-GPS融合SLAM算法梳理和代码讲解
6.彻底搞懂视觉-惯性SLAM：基于VINS-Fusion正式开课啦
7.彻底搞懂基于LOAM框架的3D激光SLAM: 源码剖析到算法优化
8.彻底剖析室内、室外激光SLAM关键算法原理、代码和实战(cartographer+LOAM +LIO-SAM)

9.从零搭建一套结构光3D重建系统[理论+源码+实践]

10.单目深度估计方法：算法梳理与代码实现

11.自动驾驶中的深度学习模型部署实战

12.相机模型与标定(单目+双目+鱼眼）

13.重磅！四旋翼飞行器：算法与实战

14.ROS2从入门到精通：理论与实战

15.国内首个3D缺陷检测教程：理论、源码与实战

16.基于Open3D的点云处理入门与实战教程

重磅！计算机视觉工坊-学习交流群已成立

扫码添加小助手微信，可申请加入3D视觉工坊-学术论文写作与投稿 微信交流群，旨在交流顶会、顶刊、SCI、EI等写作与投稿事宜。

同时也可申请加入我们的细分方向交流群，目前主要有ORB-SLAM系列源码学习、3D视觉、CV&深度学习、SLAM、三维重建、点云后处理、自动驾驶、CV入门、三维测量、VR/AR、3D人脸识别、医疗影像、缺陷检测、行人重识别、目标跟踪、视觉产品落地、视觉竞赛、车牌识别、硬件选型、深度估计、学术交流、求职交流等微信群，请扫描下面微信号加群，备注：”研究方向+学校/公司+昵称“，例如：”3D视觉 + 上海交大 + 静静“。请按照格式备注，否则不予通过。添加成功后会根据研究方向邀请进去相关微信群。原创投稿也请联系。

▲长按加微信群或投稿

▲长按关注公众号

3D视觉从入门到精通知识星球：针对3D视觉领域的视频课程（三维重建系列、三维点云系列、结构光系列、手眼标定、相机标定、激光/视觉SLAM、自动驾驶等）、知识点汇总、入门进阶学习路线、最新paper分享、疑问解答五个方面进行深耕，更有各类大厂的算法工程人员进行技术指导。与此同时，星球将联合知名企业发布3D视觉相关算法开发岗位以及项目对接信息，打造成集技术与就业为一体的铁杆粉丝聚集区，近4000星球成员为创造更好的AI世界共同进步，知识星球入口：

学习3D视觉核心技术，扫描查看介绍，3天内无条件退款

 圈里有高质量教程资料、答疑解惑、助你高效解决问题

觉得有用，麻烦给个赞和在看~