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述中提到了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天内无条件退款

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

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