浙江大学学报(工学版), 2022, 56(9): 1796-1805 doi: 10.3785/j.issn.1008-973X.2022.09.013

计算机与控制工程

融合多尺度和多头注意力的医疗图像分割方法

王万良,, 王铁军, 陈嘉诚, 尤文波

1. 浙江工业大学 计算机科学与技术学院,浙江 杭州 310023

2. 浙江树人大学 信息科技学院,浙江 杭州 310015

Medical image segmentation method combining multi-scale and multi-head attention

WANG Wan-liang,, WANG Tie-jun, CHEN Jia-cheng, YOU Wen-bo

1. College of Computer Science and Technology, Zhejiang University of Technology, Hangzhou 310023, China

2. College of Information Science and Technology, Zhejiang Shuren University, Hangzhou 310015, China

收稿日期: 2021-08-5  

基金资助: 国家自然科学基金资助项目(61873240)

Received: 2021-08-5  

Fund supported: 国家自然科学基金资助项目(61873240)

作者简介 About authors

王万良(1957—),男,教授,从事大数据、深度学习与智能调度研究.orcid.org/0000-0002-1552-5075.E-mail:zjutwwl@zjut.edu.cn , E-mail:zjutwwl@zjut.edu.cn

摘要

为了从医疗图像中自动且准确地提取兴趣区域, 提出基于神经网络的分割模型MS2Net. 针对传统卷积操作缺乏获取长距离依赖关系能力的问题, 为了更好提取上下文信息, 提出融合卷积和Transformer的架构. 基于Transformer的上下文抽取模块通过多头自注意力得到像素间相似度关系, 基于相似度关系融合各像素特征使网络拥有全局视野, 使用相对位置编码使Transformer保留输入特征图的结构信息. 为了使网络适应兴趣区域形态的差异, 在MS2Net中应用解码端多尺度特征并提出多尺度注意力机制. 对多尺度特征图依次应用分组通道和分组空间注意力, 使网络自适应地选取合理的多尺度语义信息. MS2Net在数据集ISBI 2017和CVC-ColonDB上均取得较U-Net、CE-Net、DeepLab v3+、UTNet等先进方法更优的交并比指标, 有着较好的泛化能力.

关键词: 医疗图像分割 ; 深度学习 ; 注意力 ; Transformer ; 多尺度

Abstract

A neural network based segmentation model MS2Net was proposed to automatically and accurately extract regions of interest from medical images. In order to better extract context information, a network architecture combining convolution and Transformer was proposed, which solved the problem that traditional convolution operations lacked the ability to acquire long-range dependencies. In the Transformer-based context extraction module, multi-head self-attention was used to obtain the similarity relationship between pixels. Based on the similarity relationship, the features of each pixel were fused, so that the network had a global view, while the relative positional encoding enabled Transformer to retain the structural information of an input feature map. Aiming at making the network adapt to different sizes of regions of interest, the multi-scale features of decoders were used by MS2Net and a multi-scale attention mechanism was proposed. The group channel attention and the group spatial attention were applied to a multi-scale feature map in turns, so that the reasonable multi-scale semantic information was selected adaptively by the network. MS2Net had achieved better intersection-over-union than advanced methods such as U-Net, CE-Net, DeepLab v3+, UTNet on both ISBI 2017 and CVC-ColonDB datasets, which reflected its excellent generalization ability.

Keywords: medical image segmentation ; deep learning ; attention ; Transformer ; multi-scale

PDF (1159KB) 元数据 多维度评价 相关文章 导出 EndNote| Ris| Bibtex  收藏本文

本文引用格式

王万良, 王铁军, 陈嘉诚, 尤文波. 融合多尺度和多头注意力的医疗图像分割方法. 浙江大学学报(工学版)[J], 2022, 56(9): 1796-1805 doi:10.3785/j.issn.1008-973X.2022.09.013

WANG Wan-liang, WANG Tie-jun, CHEN Jia-cheng, YOU Wen-bo. Medical image segmentation method combining multi-scale and multi-head attention. Journal of Zhejiang University(Engineering Science)[J], 2022, 56(9): 1796-1805 doi:10.3785/j.issn.1008-973X.2022.09.013

卷积神经网络(convolutional neural network, CNN)为医学图像的分割问题带来新的解决方法, 让分割和诊疗的自动化变得更加容易. 分割任务中, 全卷积网络[1](fully convolutional network, FCN)将传统CNN中的全连接层替换为卷积层以实现语义分割. 为了在获得高级语义特征的基础上保留更多的结构信息, U-Net[2]将“编码器-解码器”结构与跳跃连接结合, 使网络在解码过程中融入编码端低级特征, 具有优良的分割性能.

由于卷积操作局部感知及权值共享的特性, 基于CNN的分割模型缺乏长距离依赖关系的建模能力和灵活的空间感知能力. 不少学者研究与改进CNN的模块、结构, 旨在提高医学图像分割性能. 例如采用自然图像处理中的空间金字塔[3-5]以及非局部块[6-8]的思想来获得更好的上下文特征表示. DeepLab系列[9-10]利用空洞金字塔池化和Multi-Grid策略、Gu等[5]利用Inception结构[11]和多核池化获取更丰富的多尺度和上下文信息. Xing等[7]将可变形卷积[12]与非局部块结合, 为后续模块提供了更为可靠的特征图. Wang等[13]利用双向皮肤镜特征学习模块充分融合不同尺度的特征生成了丰富的特征表示. 学者们还从注意力的角度出发, 给特征图中的重要通道或位置自适应地赋以较大权重来突出重点并抑制噪声以提高分割性能. 通道注意力通过学习通道间的关联程度为各通道赋予对应权重, 其中“挤压−激励[14](squeeze and excitation, SE)”架构被大量用于分割任务中. 例如SFUNet[15]融合解码器各层级的输出特征图并应用SE修正特征图, 提高了眼底动静脉分割的准确度. RAUNet[16]使用基于通道注意力机制的增强注意力模块来强调编码端低级特征图的有用通道, 使得该低级特征图能与解码端的高级特征图更可靠地融合. 类似地, 空间注意力通过为各像素赋予学习到的不同权重来调整特征图. 例如Attention UNet[17]在跳跃连接处综合利用解码端和编码端的高低级特征图, 生成关于低级特征图每个像素的门控信息. 不少研究注重联合多种注意力来提高任务性能. 例如BAM[18]和CBAM[19]在通道和空间维度上独立地精炼卷积特征, CA-Net[20]全面结合各种注意力机制以提高分割任务的准确性和解释性. 此外, 自注意力可以有效克服上下文信息提取能力有限的问题. 例如Liu等[8]利用非局部块的思想融合全局信息. Transformer在自然语言处理(natural language processing, NLP)中取得巨大成功, 其中的多头自注意力机制通过上下文聚合算子实现成对的元素交互, 具有捕获远程关联特征的能力. 学者们先后提出ViT[21]、Swin Transformer[22]、CMT[23]等来解决视觉任务问题, 在分割任务上Cao等[24]提出基于纯Transformer结构的分割网络Swin-Unet, Gao等[25]结合Transformer和卷积多方面的优势和不足, 提出UTNet, 为解决医学图像分割问题扩展了思路.

本研究综合考虑上下文和注意力的关键内容和技术, 提出融合多尺度和多头注意力的分割网络框架(MS2Net). MS2Net由特征提取模块、上下文抽取模块(context extractor, CE)、多尺度注意力(multi-scale attention, MSA)和特征解码模块组成, 旨在提升U-Net在医疗图像上的分割效果. 本研究将在ISBI 2017[26]和CVC-ColonDB[27]数据集上测试MS2Net的分割性能.

1. MS2Net

图1所示, MS2Net以分割网络U-Net的“编码器-解码器”结构作为骨干. 其中编码器被视为收缩路径上获取各尺度特征的特征提取器, 各尺度的编码器组成特征提取模块; 解码器被视为扩张路径上将高级语义特征恢复成更高分辨率特征图的上采样器, 各尺度的解码器组成特征解码模块. 在每个尺度上, 编码器输出端通过长跳跃连接结构连接到对应尺度的解码器上, 以融合该尺度上的低级特征图(编码器端)和高级特征图(解码器端)的信息, 再将处理后的融合特征图作为下个解码器的输入. 在该尺度上的最终解码输出, 既有高级的语义信息, 又有低级的空间结构信息, 提高了分割网络的准确性和鲁棒性. 传统卷积操作的核大小固定, 传统卷积操作一般仅从邻域像素收集信息, 缺乏显式捕获远距离依赖关系的能力. MS2Net基于Transformer, 利用多头自注意力机制获取长距离依赖关系, 提出上下文抽取模块, 并利用位置编码机制保留输入特征图的结构信息. 为了适应医学图像兴趣区域差别较大的分割任务, MS2Net在解码端拼接多个尺度上的特征图, 利用多尺度注意力, 学习不同尺度特征图间的关系和各尺度的特点, 突出各尺度上与分割目标形态大小相关的特征.

图 1

图 1   MS2Net框架图与局部模块放大图

Fig.1   MS2Net’s frame diagram and local modules’ details


1.1. 特征提取模块和特征解码模块

特征提取模块用于提取输入医学图像的高层语义信息, 通过CNN连续的卷积和池化操作逐步降低特征图的分辨率, 产生有关图像局部区域的抽象特征表示. U-Net每个编码块包含2个卷积层和1个最大池化层. MS2Net采用ResNet34[28]作为特征提取模块, 将ResNet起始步长为2的7×7卷积替换为2个连续步长为1的3×3卷积, 使得特征提取模块整体的下采样率为16. 结合迁移学习方法, 本研究使用在大规模数据集ImageNet上预训练的模型初始化特征提取模块中的部分参数. 与随机初始化相比, 迁移学习使网络能够更快地找到最优解, 提高模型的收敛速度.

特征解码模块旨在将低分辨率的特征图恢复成高分辨率特征图, 并融合对应尺度的编码端低级特征图. 本研究采用1×1卷积、3×3转置卷积、1×1卷积的串行结构来恢复特征图. 在各尺度编码端和解码端输出特征图的融合策略上, 区别于U-Net“跳连+堆叠”融合高低级信息的形式, MS2Net采用将同个尺度上的编码端低级特征图和解码端恢复的高级特征图直接相加的形式.

1.2. 上下文抽取模块

上下文抽取模块用于融合输入特征图的上下文信息, 传统的卷积操作仅从邻域像素收集信息, 缺乏显式捕获长距离依赖关系的能力. 新算子(如空洞卷积、可变形卷积)的出现为更有效地扩大感受野、更好地提取上下文信息提供了新途径, 也在一定程度上提高了分割的性能, 但是固定的卷积核大小仍然限制了CNN对于输入内容的接收能力. NLP中基于自注意力的Transformer架构较为成熟, 其中具有强大提取长距离依赖关系能力的自注意力, 通过分析序列中的相似关系来聚合相关特征. 本研究将数个串联的带位置感知的Transformer块作为上下文抽取模块. 在每个Transformer块中, 先使用卷积提取输入特征图的局部信息, 为了减少计算参数量, 采用3×3的深度可分离卷积和1×1卷积, 并利用跳连结构缓解可能出现的梯度消失状况.

1.2.1. 多头自注意力

多头自注意力从各子空间分别推断像素间相关性, 得到具有全局视野的特征图, 通过堆叠不同自注意力头对应的结果特征图, 使语义表示更为丰富.

假设注意力头数为N, 令输入特征图为 ${\boldsymbol{X}} \in {{{\bf{R}}}^{D \times h \times w}}$, 其中 $ D $为通道个数, hw分别为特征图的高和宽. 通过维度调整和线性变换使 ${\boldsymbol{X}}$变换为3个矩阵, 分别为查询矩阵 ${\boldsymbol{Q}} \in {{{\bf{R}}}^{hw \times Nk}}$、键矩阵 ${\boldsymbol{K}} \in {{{\bf{R}}}^{hw \times Nk}}$、值矩阵 ${\boldsymbol{V}} \in {{{\bf{R}}}^{hw \times Nv}}$, 其中, $ k $${\boldsymbol{K}}$${\boldsymbol{Q}}$的列数的1 / N, $ v $${\boldsymbol{V}}$的列数的1 / N. 将 ${\boldsymbol{Q}}$${\boldsymbol{K}}$${\boldsymbol{V}}$按列划分成N份, 每份分别记为 ${{\boldsymbol{Q}}^n} \in {{{\bf{R}}}^{hw \times k}}$${{\boldsymbol{K}}^n} \in {{{\bf{R}}}^{hw \times k}}$${{\boldsymbol{V}}^n} \in {{{\bf{R}}}^{hw \times v}}$, $ n \in [1,N] $, 分别成为第n个注意力头的查询矩阵、键矩阵和值矩阵. 第n个输出头 ${{\boldsymbol{O}}^n}$

$ {{\boldsymbol{O}}^n} = {\text{SA}}({{\boldsymbol{Q}}^n},{{\boldsymbol{K}}^n},{{\boldsymbol{V}}^n}) = {\text{SoftMax}}\left(\dfrac{{{{\boldsymbol{Q}}^n} {{({{\boldsymbol{K}}^n})}^{\rm{T}}}}}{{\sqrt k }}\right){{\boldsymbol{V}}^n} . $

式中: SA为自注意力函数. 式子由2个部分组成, 分别是相似度矩阵计算和输出头计算. 其中, 相似度矩阵由SoftMax归一化得到, 输出头由相似度矩阵和值矩阵相乘得到.

堆叠各个输出头的特征图, 生成多头自注意力的输出为

$ {\text{MHSA}}({\boldsymbol{Q}},{\boldsymbol{K}},{\boldsymbol{V}}) = {\text{Concat}}([ {{{\boldsymbol{O}}^1}},\; \cdots ,\;{{{\boldsymbol{O}}^n}},\; \cdots ,\;{{{\boldsymbol{O}}^N}} ]). $

式中: MHSA为多头自注意力函数, Concat为堆叠函数.

在多头自注意力中, 各元素间的内积操作使自注意力在本质上具有全局感受野, 并且擅长捕捉长距离依赖. 由于自注意力机制仅仅计算不同像素特征间的相关性, 没有考虑各像素的位置信息, 缺乏空间感知能力, 得到的输出会损失结构信息, NLP中常用的绝对位置编码在计算SoftMax时会损失重要的相对位置关系, 为此本研究使用相对位置编码(relative positional encoding, RPE)保留结构信息.

1.2.2. RPE

基于Liu等[22]的方法, 为每个注意力头设置 $ (2h - 1) \times (2w - 1) $可训练的相对位置信息矩阵 $ {{\boldsymbol{P}}^n} $, 用以表示 $ (2h - 1) \times (2w - 1) $个相对位置关系, 并改写式(1), 使之融入相对位置信息:

$ {{\boldsymbol{O}}^n} = {\text{SA}}({{\boldsymbol{Q}}^n},{{\boldsymbol{K}}^n},{{\boldsymbol{V}}^n}) = {\text{SoftMax}}\left(\frac{{{{\boldsymbol{Q}}^n} {{({{\boldsymbol{K}}^n})}^{\rm{T}}}}}{{\sqrt k }}+{{\boldsymbol{B}}^n}\right){{\boldsymbol{V}}^n} . $

式中: ${\boldsymbol{B}}^n $为各个像素间的位置相关性, ${{\boldsymbol{B}}^n} \in {{{\bf{R}}}^{hw \times hw}}$; $ B_{i,j}^n $为位置j相对于位置i的第n个注意力头对应的相对位置信息, 代表着位置j相对于位置i在空间上的重要程度, $ B_{i,j}^n $可以通过查相对位置矩阵得到. 假设位置j对应的像素坐标为 $ ({x_j},{y_j}) $, 位置i对应的像素坐标为 $ ({x_i},{y_i}) $, 那么 $ B_{i,j}^n = P_{h+{y_j} - {y_i},w+{x_j} - {x_i}}^n $. ${{\boldsymbol{P}}^n}$可以通过梯度反向传播来学习各相对位置的权重值.

融合了相对位置信息, 多头注意力机制能够在捕捉长距离依赖关系的同时, 保留输入Transformer的特征图中的结构信息.

1.2.3. 前馈网络

在Transformer中, 前馈网络一般采用2个线性层. 第1个线性层将维度扩充为原来的4倍, 采用GeLU等函数作为激活函数, 第2个线性层用来恢复维度. 结合Guo等[23]的设计, 本研究采用的前馈网络由1×1卷积、3×3深度可分离卷积、1×1卷积组成, 在每个卷积层后都进行批归一化操作, 并将GeLU作为激活函数; 在深度可分离卷积处引入跳连结构, 以促进梯度的跨层传播能力.

1.3. 多尺度注意力

“编码器−解码器”结构包含多个不同的尺度, 利用多个尺度的特征图学习不同尺度间关系, 突出各个尺度上与分割目标相关的特征将有利于提升分割性能. 基于Woo等[19-20]的方法, 本研究利用多尺度注意力自动学习每个尺度特征图的空间权重和通道权重, 使网络具备适应不同大小病灶的泛化能力, 其框架结构如图2所示. 将特征解码模块(扩张路径)各尺度上的特征图分别进行卷积和上采样操作以统一大小. 以图1中特征解码模块的4个尺度为例, 每个尺度经过卷积、上采样操作后都得到 $ m \times H \times W $的特征图, 其中m为每个尺度对应的处理后的特征图通道数. 每个尺度的特征图记为 $ {{\boldsymbol{S}}_i}(i = 1,2,3,4) $, 堆叠 $ {{\boldsymbol{S}}_i} $形成张量 $ {\boldsymbol{S}} \in {{{\bf{R}}}^{4m \times H \times W}} $.

图 2

图 2   多尺度注意力机制

Fig.2   Multi-scale attention mechanism


1.3.1. 分组通道注意力

多尺度注意力须关注各尺度间通道的关系, 为此使用“挤压−激励”调整不同尺度的通道权重. 为了更好地体现多尺度的特性, 增加分组通道融合策略使各尺度共享相同的通道权重, 形成分组通道注意力(group channel attention, GCA). 具体操作如下.

1) 输入特征图 $ {\boldsymbol{S}} $, 通过全局平均池化(global average pooling, GAP)操作 $ {F_{{\text{sq}}}}( \cdot ) $, 将 $ {\boldsymbol{S}} $压缩成长度为4m的向量z, z中的元素 $ {z_c} $包含通道c的全局信息, 即

$ {z_c} = {F_{{\text{sq}}}}({{\boldsymbol{S}}_c}) = \frac{1}{{H \times W}}\sum\limits_{i = 1}^H {\sum\limits_{j = 1}^W {{S_{c,i,j}}} } . $

式中: $ {{\boldsymbol{S}}_c} $$ {\boldsymbol{S}} $c个通道的特征.

2) 采用多层感知机建立通道相关性模型. 多层感知机由2个全连接层组成: (a)全连接层输出通道数为 $ 4m/r $, 其中 $ r $为缩放因子, 取 $ r = 4 $, 采用ReLU作为激活函数; (b)全连接层输出通道数为4m. 最终得到各通道的权重 $ {\boldsymbol{\omega }} $, 即

$ {\boldsymbol{\omega }} = {{\boldsymbol{\varOmega }}_2} {\rm{ReLU}} \; ({{\boldsymbol{\varOmega }}_1} {\boldsymbol{z}}) . $

式中: ${{\boldsymbol{\varOmega }}_1}$${{\boldsymbol{\varOmega }}_2}$分别为全连接层(a)和(b)的参数矩阵, 维度分别为 $ \left( {4m/r} \right) \times 4m $$ 4m \times \left( {4m/r} \right) $.

3) 利用分组通道融合策略对各尺度对应的权重进行加权平均, 将加权后的均值作为该尺度对应的所有通道的权重, 即

$ {\overline \omega _c} = \frac{{\displaystyle\sum\limits_{i = \left\lfloor {\left( {c - 1} \right)/m} \right\rfloor \times m+1}^{\left\lfloor {\left( {c - 1} \right)/m} \right\rfloor \times m+m} {\exp\; ({\alpha _i}){\omega _i}} }}{{\displaystyle\sum\limits_{i = \left\lfloor {\left( {c - 1} \right)/m} \right\rfloor \times m+1}^{\left\lfloor {\left( {c - 1} \right)/m} \right\rfloor \times m+m} {\exp \;({\alpha _i})} }},\;c \in [1,4m] \cap \bf{Z} . $

式中: $ \lfloor ·\rfloor $为下取整函数, $ {\alpha _i} $为可训练向量 $ {\boldsymbol{\alpha }} $的第i个元素.

4) 通过激活函数Sigmoid将权重 $ {\boldsymbol{\overline \omega }} $归一化到 (0, 1.0), 得到激活向量 $ {{\boldsymbol{\gamma }}^{{\text{CA}}}} $. 将产生该激活向量的过程记为 ${F_{{\text{ex}}}}( \cdot ,{\boldsymbol{\varOmega }})$, 即

$ {{\boldsymbol{\gamma }}^{{\text{CA}}}} = {F_{{\text{ex}}}}({\boldsymbol{z}},{\boldsymbol{\varOmega }}) = \sigma ({\boldsymbol{\overline \omega }}) = \frac{1}{{1+{\exp \; ({ - {\boldsymbol{\overline \omega }}})}}} . $

5) 利用 $ {{\boldsymbol{\gamma }}^{{\text{CA}}}} $加权 $ {\boldsymbol{S}} $每个特征通道, 得到的特征图 $ {{\boldsymbol{S}}^{{\text{CA}}}} = {\boldsymbol{S}} \otimes {{\boldsymbol{\gamma }}^{{\text{CA}}}} $. 其中 $ \otimes $为带广播的元素乘法.

1.3.2. 分组空间注意力

分组空间注意力(group spatial attention, GSA)将 $ {{\boldsymbol{S}}^{{\text{CA}}}} $先后通过3×3卷积+ReLU和1×1卷积(4通道输出)+Sigmoid来获得不同尺度的空间权重 ${{\boldsymbol{\gamma }}^{{\text{SA}}}} \in {{\bf{R}}^{4 \times H \times W}}$. 这样得到的各尺度的权重不仅考虑了不同位置的关系, 也融合了不同尺度的信息. 得到的特征图满足

$ {\boldsymbol{S}}_c^{{\text{SA}}} = {\boldsymbol{S}}_c^{{\text{CA}}} \otimes {\boldsymbol{\gamma }}_{\left\lfloor {\frac{{c - 1}}{m}} \right\rfloor +1}^{{\text{SA}}},\;c \in [1,4m] \cap \bf{Z} . $

本研究的多尺度注意力模块的最终输出为

$ {{\boldsymbol{S}}^{\text{O}}} = {{\boldsymbol{S}}^{{\text{SA}}}}+{{\boldsymbol{S}}^{{\text{CA}}}}+{\boldsymbol{S}} . $

在计算多尺度注意力模块输出时用到的残差连接可以促进梯度的反向传播, 加快网络的训练过程. 通过多尺度注意力机制, 网络可以自适应地从各尺度获取有用的信息, 特别适合医学图像兴趣区域大小差异较大的场景.

1.4. 损失函数

本研究的框架是端到端可学习的网络, 方法面向医疗图像(比如皮肤镜影像)的兴趣区域分割, 即区分每个像素是前景区域还是背景区域. 为了更好地突出对小目标的关注程度, 采用Dice函数作为损失函数进行优化. Dice损失函数LDice与交并比度量直接相关, 计算式为

$ {L_{{\text{Dice}}}} = 1 - \displaystyle\sum\limits_c^{} {\frac{{2{\eta _c}\displaystyle\sum\limits_{i = 1}^H {\displaystyle\sum\limits_{j = 1}^W {{p_{c,i,j}}{g_{c,i,j}}} } }}{{\displaystyle\sum\limits_{i = 1}^H {\displaystyle\sum\limits_{j = 1}^W {{p_{c,i,j}}} } +\displaystyle\sum\limits_{i = 1}^H {\displaystyle\sum\limits_{j = 1}^W {{g_{c,i,j}}} } }}} . $

式中: $ {p_{c,i,j}} $为在位置 $ (i,j) $处的像素被预测为类别c的概率; $ {g_{c,i,j}} $用于表示真实标签, 当 $ {g_{c,i,j}} = 1 $时,表明当前像素输入类别c, 否则不属于类别c$ {\eta _c} $为类别c的权重, 满足 $\displaystyle\sum\nolimits_c {{\eta _c}} = 1$.

2. 实验结果与分析

2.1. 数据集和评价指标

使用皮肤镜影像数据集ISBI 2017和无线胶囊内窥镜数据集CVC-ColonDB训练并评估MS2Net. ISBI 2017数据集包含3个部分, 分别是训练集、验证集和测试集, 其中训练集包含2 000张图片, 验证集包含150张图片, 测试集包含600张图片. CVC-ColonDB包含380张内窥镜影像图.

评估病灶区域分割性能的指标包括Jaccard相似度指数(JA)、Dice相似系数(DI)、准确度(ACC)、敏感度(SEN)、特异度(SPE)、带阈值的Jaccard指数(TJI)来, 计算式分别为

$ {\text{JA}} = \frac{{\left| {P \cap G} \right|}}{{\left| {P \cup G} \right|}} , $

$ {\text{DI}} = \frac{{2\left| {P \cap G} \right|}}{{\left| P \right|+\left| G \right|}} , $

$ {\text{ACC = }}\frac{{{\text{TP+TN}}}}{{{\text{TP+TN+FP+FN}}}} , $

$ {\text{SEN = }}\frac{{{\text{TP}}}}{{{\text{TP+FN}}}} , $

$ {\text{SPE = }}\frac{{{\text{TN}}}}{{{\text{TN+FP}}}} , $

$ {\text{TJI}} = \left\{ \begin{aligned} & {{\text{JSI}},\,\,\;\; \,{\text{JSI}} \geqslant 0.65}; \\ & {0,\,\,\,\,\,\,\;\;\;\;{\text{JSI}} < 0.65} . \end{aligned} \right. $

式中: P为预测标签, G为真实的标签, TP为预测正确且预测为正类的像素数量, TN为预测正确且预测为负类的像素数量, FP为预测错误且预测为正类的像素数量, FN为预测错误且预测为负类的像素数量. JA、DI指标侧重小目标的评价, 分割目标越小, 其越敏感; ACC、SEN和SPE用来度量分割的准确性, 评估的目的是判断相关像素是否被正确分类; TJI是JA的带阈值版本, 参考文献[29]将最小阈值设置为0.65.

2.2. 实验设置

为了在提高训练速度的同时保证算法对比的公平性(如批大小相同),MS2Net和对比算法均使用PyTorch框架实现, 在多张NVIDIA GeForce RTX 2080 Ti显卡上并行化训练. 所有方法均使用相同的损失函数(Dice损失函数)和Adam优化器训练. 实验中训练采用Poly策略, 学习率初始值为0.001, 计算式为

$ {l_e} = {l_{\text{0}}} {\left(1 - \frac{e}{E}\right)^{\tau}} . $

式中: $ E $为最大训练周期数; $ e $为训练过程中的某个周期, $ e \in [0,E - 1] $$ {l_{\text{0}}} $为学习率初始值; $ {l_e} $为第 $ e $个周期对应的学习率; $ \tau $为衰减系数,设置为0.9. 对于ISBI 2017数据集, 设置所有方法的最大训练周期数均为300, 训练阶段的批大小均为40; 对于CVC-ColonDB,设置最大训练周期为500, 批大小为16.

ISBI 2017数据集中的图像尺寸不一, 在训练集中高×宽的范围为542×718至4499×6748, 宽高的比例近似4∶3, 因此将图像缩放到高为216、宽为288. 在训练过程中, 将192×256的随机裁剪、按概率水平和垂直翻转以及介于(−π/6, π/6)的角度随机旋转作为数据增强手段. 在测试推理阶段, 采用缩放后的192×256图片评估网络的性能. CVC-ColonDB数据集中的每张图像的尺寸均为500×574, 实验过程中将其缩放到224×256作为网络输入, 在训练阶段采用随机旋转、翻转方法作为数据增强的手段, 实验中采用五折交叉验证的方法评估网络性能.

2.3. 网络设计评估

采用消融研究方法验证各模块的有效性, ISBI 2017测试集上的表现如表1所示. 在比较中, 实验设置上下文抽取模块Transformer块数T=4, 注意力头数N=32个, 并将特征提取模块和特征解码模块的组合称为基准网络BN. 可以看到, 添加上下文抽取模块后, JA、DI、ACC指标较BN分别提升0.64%、0.29%和0.12%, 采用多尺度注意力机制后,这些指标分别提升0.81%、0.48%和0.32%. 表明所提出的上下文抽取模块和多尺度注意力机制均能够充分提取特征.多尺度注意力通过充分结合各尺度的关系与特征, 能够适应医学图像中兴趣区域的差异, 有利于医学图像分割. 综合使用上下文抽取模块和多尺度注意力, 将给网络带来更好的性能, 例如MS2Net较BN在JA指标上提高1.16 %, 达到78.43 %, 在准确度、敏感度上也有优秀的表现.

表 1   不同模块组合方式对分割性能的影响

Tab.1  Effects of different module combinations on segmentation performance

%
组合方式 JA DI ACC SEN SPE TJI
BN 77.27 85.71 93.65 84.70 96.27 68.26
BN+MSA 78.08 86.19 93.97 85.95 96.17 69.85
BN+CE 77.91 86.00 93.77 84.63 96.67 70.01
MS2Net 78.43 86.28 93.81 85.96 96.72 71.60

新窗口打开| 下载CSV


2.3.1. 特征提取和解码模块分析

MS2Net的特征提取模块主要是基于预训练的ResNet34设计的, 其中的预训练模型是在大型数据集ImageNet上训练得到的. 为了评估预训练对分割性能的影响, 本研究在相同条件(N=32, T=4)下进行对比实验, 实验结果如表2所示. 在各个评价指标上, 使用预训练特征提取模块的效果较未使用的模型有明显的提高. 以JA指标为例, 使用预训练特征提取模块参数的模型较未使用预训练参数的网络高0.46%. 可见, 预训练模型有利于网络获得更优的性能, 让网络跳出可能的局部最优点或鞍点.

表 2   预训练对分割性能的影响

Tab.2  Effect of pre-training on segmentation performance

%
预训练 JA DI ACC SEN
77.97 85.98 93.45 85.73
78.43 86.28 93.81 85.96

新窗口打开| 下载CSV


传统的U-Net结构在特征解码阶段, 采用编码端低级特征图和解码端高级特征图堆叠的方式来融合信息, 使得解码器在获得高层特征的同时, 获得下采样过程中丢失的结构特征. 如表3所示,采用相加的融合方式比采用堆叠的融合方式更加可靠, 在N=32、T=4的情况下, 前者比后者在JA指标上高0.60%, 并且引入更少的参数量NP. 比较两者的浮点计算数(floating point operations, FLOPs)可以看出, 相加的融合方式引入的计算复杂度比堆叠的融合方式小.

表 3   跳跃连接时特征图融合方式对分割性能的影响

Tab.3  Influence of feature map fusion mode on segmentation performance during skip connection

方式 JA/% DI/% ACC/% SEN/% NP/106 FLOPs/109
堆叠 77.83 85.96 93.79 84.55 22.47 17.59
相加 78.43 86.28 93.81 85.96 21.66 16.96

新窗口打开| 下载CSV


2.3.2. 上下文抽取模块分析

表4所示,设置注意力头数N=32, 分别观察Transformer块数T=1、2、4、8时网络的表现.可以看到, 在一定范围内, 随着块数的增加, 上下文抽取模块融合上下文信息更加充分. 总体来看, 当T=4时, 各项指标的综合表现最好, 且在JA、DI、ACC这3项指标中取得最优值.

表 4   注意力头数为32时不同Transformer块数对分割指标的影响

Tab.4  Influence of different Transformer block numbers on segmentation indexes when number of heads is thirty-two

%
T JA DI ACC SEN
1 77.97 86.03 93.80 85.42
2 78.08 86.19 93.67 85.56
4 78.43 86.28 93.81 85.96
8 78.19 86.17 93.80 86.18

新窗口打开| 下载CSV


设置T=4, 研究不同头数对分割性能的影响, 分别选取N=1、2、4、8、16、32进行比较实验, 实验结果如表5所示. 在交并比相关指标方面, 随着自注意力头数的增加, JA、DI整体呈上升趋势, 当N=32时, JA、DI指标分别达到78.43%和86.28%, 为6组中的最优值; 在准确度方面, ACC=93.81%, 有较高准确度. 结果表明, 在合理范围内, 较多的自注意力头将有利于提升分割性能. 原因是多头的注意力有助于网络捕捉多个模式的信息, 通过综合利用各方面的信息, 网络将更具鲁棒性.

表 5   Transformer块数为4时不同注意力头数对分割指标的影响

Tab.5  Influence of different number of heads on segmentation indexes when number of Transformer block is four

%
N JA DI ACC SEN
1 77.53 85.69 93.65 85.42
2 77.75 85.91 93.58 85.33
4 77.84 85.91 93.83 84.97
8 77.88 85.85 93.24 86.24
16 78.11 86.22 93.93 85.99
32 78.43 86.28 93.81 85.96

新窗口打开| 下载CSV


为了观察相对位置编码的作用, 在N=32、T=4的情况下, 对有无相对编码的情形分别进行网络训练与模型评估, 如表6所示. 可以看到, 在各项指标上, 有相对编码的效果更胜一筹. 引入相对位置编码, 在两两像素特征间相似程度计算的同时, 考虑了两点距离对相似度的影响, 由此自注意力也在拥有全局视野的同时, 拥有了对相邻重点元素信息的感知和提取能力.

表 6   相对编码对分割性能的影响

Tab.6  Effect of relative positional encoding on segmentation performance

%
相对位置编码 JA DI ACC SEN
不采用 77.86 85.94 93.76 84.50
采用 78.43 86.28 93.81 85.96

新窗口打开| 下载CSV


2.3.3. 多尺度注意力机制分析

为了评价分组通道注意力和分组空间注意力在多尺度注意力中的作用, 在T=4、N=32的情况下进行实验对比, 实验结果如表7所示. 结果表明, 相比直接将各层次特征图堆叠起来进行后续像素级别分类的方法, 利用单一分组注意力虽然使ACC下降, 但是其分割交并比指标JA、DI分别上升0.38%和0.13%. 综合使用2种分组注意力相较于直接融合方法在JA、DI、ACC指标上分别提高0.68%、0.41%和0.16%.

表 7   分组通道注意力和分组空间注意力在分割指标上的作用对比

Tab.7  Contrast between group channel attention and group spatial attention on segmentation index

%
GCA GSA JA DI ACC SEN
77.75 85.87 93.65 84.53
78.13 86.15 93.54 85.72
77.88 85.94 93.49 85.89
78.43 86.28 93.81 85.96

新窗口打开| 下载CSV


图3所示为不同病灶的图像在不同尺度上的空间注意力、通道注意力可视化效果及相关数值. 图中, G为分组通道权重. 可以看到, 在不同尺度上注意力有不同的侧重点. 在病灶大小不同的病历中, 可以观察到在尺度2和尺度3上, 大病灶特征图的分组通道注意力系数更大, 重要性更高; 在尺度1上,小病灶的特征图的重要性更大. 这表明, 多尺度注意力能够自动适应不同形状大小的病灶, 为不同形态病灶的各个尺度特征图自适应地注入特异的关注度, 进而更好分割不同大小的病变. 还可以看到, 在尺度较大的特征图(如尺度1、2)中, 多尺度注意力模块倾向于注意病灶的边缘, 以获得清晰的病灶边界.

图 3

图 3   不同病灶图像的注意力权重可视化图

Fig.3   Attention weight visualization of lesions of different scales


分别利用直接融合、SE、CBAM和MS-Dual-Guided[30]等方式对各层级解码端输出的特征图进行融合, 并与多尺度注意力进行比较分析, 结果如表8所示. 结果表明, 多尺度注意力相较于其他融合策略有更优的表现.

表 8   不同融合方式的各项分割指标对比

Tab.8  Comparison of segmentation indexes of different fusion methods

%
方式 JA DI ACC SEN
直接融合 77.75 85.87 93.65 84.53
SE 77.59 85.79 93.48 85.25
CBAM 78.02 86.03 93.62 85.74
MS-Dual-Guided 77.83 85.80 93.60 85.86
MSA 78.43 86.28 93.81 85.96

新窗口打开| 下载CSV


2.4. 对比实验

将MS2Net与其他“编码器-解码器”架构的医学图像分割算法进行比较, 除了网络框架不同外, 数据处理模式、数据增强方式、损失函数以及训练机制均相同, 各方法在ISBI 2017数据集上的结果如表9所示. 可以看到, 相较于使用变种卷积和池化来提高上下文信息的CE-Net, MS2Net在交并比指标上提高0.97%. 相较于基于Transformer的医疗图像分割网络Swin-UNet, MS2Net取得超过12%的交并比优势. 在大尺度特征图的计算问题上, Swin Transformer采用基于滑动窗口的多头注意力机制调整窗口内的像素, 而不同窗口间除了在前馈网络中有半窗口滑动外缺乏信息交互, 因此在数据量较小的情况下, 分割效果不佳且可能有明显的分块. 与基于融合策略的UTNet相比, MS2Net在JA指标上提高了0.96%, 且在ACC上也有一定程度的提高. 在参数量和计算量方面, MS2Net与对比算法的平均水平相当.

表 9   不同算法在ISBI 2017数据集上的分割性能对比

Tab.9  Comparison of segmentation performance of different algorithms on ISBI 2017 dataset

%
方法 JA DI ACC SEN SPE TJI
U-Net 72.81 81.78 92.23 80.36 97.33 60.83
Attention U-Net 72.93 81.89 92.10 81.72 96.97 61.27
Swin-Unet 66.04 75.61 90.46 79.11 93.81 52.58
RAUNet 77.26 85.49 93.68 83.48 97.50 69.47
SFUNet 76.15 84.57 93.38 82.98 96.83 67.01
DeepLab v3+
(Xception)
77.37 85.67 93.61 83.96 96.90 69.10
CE-Net 77.46 85.43 93.68 83.84 97.12 70.49
CA-Net 77.16 85.38 93.13 85.80 95.53 68.56
UTNet 77.47 85.51 93.55 87.15 95.48 70.10
MS-Dual-Guided 76.48 84.72 92.65 87.11 94.54 68.17
MS2Net 78.43 86.28 93.81 85.96 96.72 71.60

新窗口打开| 下载CSV


在CVC-ColonDB数据集上, 采用五折交叉验证的模型评估结果如表10所示. 表中数据项均为交叉验证的均值和标准差形式. 结果表明, MS2Net取得了优秀的分割性能, JA、DI指标高于医疗分割领域的优秀算法CE-Net、CA-Net、UTNet, 在分割准确度方面MS2Net也有优异表现. 综合上述实验, 可以验证MS2Net具有较强的可靠性和泛化能力, 适于在多种医疗图像分割任务中使用.

表 10   不同算法在CVC-ColonDB数据集上的分割性能对比

Tab.10  Comparison of segmentation performance of different algorithms on CVC-ColonDB dataset

%
方法 JA DI ACC SEN
U-Net 76.70±1.73 83.97±1.71 97.85±0.19 84.12±3.10
Attention U-Net 76.71±3.00 83.65±2.95 97.98±0.40 83.61±3.35
Swin-Unet 34.36±2.62 44.84±2.92 92.05±0.65 54.79±2.43
RAUNet 82.41±1.79 88.81±1.89 98.63±0.20 89.12±1.71
SFUNet 80.12±0.79 87.23±0.81 98.32±0.13 87.86±1.77
DeepLab v3+
(Xception)
79.34±1.24 85.61±1.41 98.52±0.10 86.19±1.38
CE-Net 81.71±1.65 87.95±1.77 98.50±0.08 88.79±2.08
CA-Net 76.01±1.60 83.73±1.63 97.57±0.29 85.35±1.93
UTNet 78.39±2.32 85.51±2.30 98.20±0.21 85.96±2.68
MS2Net 82.83±1.71 89.19±1.89 98.52±0.22 89.68±1.95

新窗口打开| 下载CSV


图4所示, 选取典型的分割结果对比不同方法的分割能力. 图中, 前4行为ISBI 2017测试集的结果, 后3行是CVC-ColonDB数据集在同一次交叉验证下的测试集的结果. 可见有了上下文抽取模块和多尺度注意力机制, MS2Net可以更加可靠地提取有用信息, 且能更好地适应数据集兴趣区域差异较大的特点.

图 4

图 4   不同方法的医疗图像分割结果对比

Fig.4   Comparison of medical image segmentation results of different methods


3. 结 语

上下文信息和多尺度信息是影响医疗图像分割性能的关键因素. 本研究提出的融合多尺度和多头注意力的分割网络能够有效地结合卷积和Transformer的优势, 在保有强大的局部特征提取能力的同时, 有效地提取长距离依赖关系. 本研究所提方法综合考虑解码端的多尺度语义特征, 使网络拥有更强的鲁棒性, 能够适应病灶大小变化较大的场景. 在接下来的工作中, 将研究轻量级的融合框架, 在保证分割性能的同时, 有效地降低参数量和计算复杂度.

参考文献

LONG J, SHELHAMER E, DARRELL T. Fully convolutional networks for semantic segmentation [C]// Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. Boston: IEEE, 2015: 3431-3440.

[本文引用: 1]

RONNEBERGER O, FISCHER P, BROX T. U-Net: convolutional networks for biomedical image segmentation [C]// International Conference on Medical Image Computing and Computer-Assisted Intervention. Munich: Springer, 2015: 234-241.

[本文引用: 1]

HE K, ZHANG X, REN S, et al

Spatial pyramid pooling in deep convolutional networks for visual recognition

[J]. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2014, 37 (9): 1904- 1916

[本文引用: 1]

ZHAO H, SHI J, QI X, et al. Pyramid scene parsing network [C]// Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. Honolulu: IEEE, 2017: 2881-2890.

GU Z, CHENG J, FU H, et al

CE-Net: context encoder network for 2D medical image segmentation

[J]. IEEE Transactions on Medical Imaging, 2019, 38 (10): 2281- 2292

DOI:10.1109/TMI.2019.2903562      [本文引用: 2]

WANG X, GIRSHICK R, GUPTA A, et al. Non-local neural networks [C]// Proceedings of the IEEE conference on computer vision and pattern recognition. Salt Lake City: IEEE, 2018: 7794-7803.

[本文引用: 1]

XING X, YUAN Y, MENG M Q

Zoom in lesions for better diagnosis: attention guided deformation network for WCE image classification

[J]. IEEE Transactions on Medical Imaging, 2020, 39 (12): 4047- 4059

DOI:10.1109/TMI.2020.3010102      [本文引用: 1]

LIU R, LIU M, SHENG B, et al

NHBS-Net: a feature fusion attention network for ultrasound neonatal hip bone segmentation

[J]. IEEE Transactions on Medical Imaging, 2021, 40 (12): 3446- 3458

DOI:10.1109/TMI.2021.3087857      [本文引用: 2]

CHEN L C, PAPANDREOU G, KOKKINOS I, et al

DeepLab: semantic image segmentation with deep convolutional nets, atrous convolution, and fully connected CRFs

[J]. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2017, 40 (4): 834- 848

[本文引用: 1]

CHEN L C, PAPANDREOU G, SCHROFF F, et al. Rethinking atrous convolution for semantic image segmentation [EB/OL]. [2021-08-01]. https://arxiv.org/abs/1706.05587.

[本文引用: 1]

SZEGEDY C, IOFFE S, VANHOUCKE V, et al. Inception-V4, inception-ResNet and the impact of residual connections on learning [C]// Thirty-first AAAI Conference on Artificial Intelligence. San Francisco: AAAI, 2017: 4278-4284.

[本文引用: 1]

ZHU X, HU H, LIN S, et al. Deformable convNets V2: more deformable, better results [C]// Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. Long Beach: IEEE, 2019: 9308-9316.

[本文引用: 1]

WANG X, JIANG X, DING H, et al

Bi-directional dermoscopic feature learning and multi-scale consistent decision fusion for skin lesion segmentation

[J]. IEEE Transactions on Image Processing, 2019, 29: 3039- 3051

[本文引用: 1]

HU J, SHEN L, SUN G. Squeeze-and-excitation networks [C]// Proceedings of the IEEE conference on computer vision and pattern recognition. Salt Lake City: IEEE, 2018: 7132-7141.

[本文引用: 1]

高颖琪, 郭松, 李宁, 等

语义融合眼底图像动静脉分类方法

[J]. 中国图象图形学报, 2020, 25 (10): 2259- 2270

DOI:10.11834/jig.200187      [本文引用: 1]

GAO Ying-qi, GUO Song, LI Ning, et al

Arteriovenous classification method in fundus images based on semantic fusion

[J]. Journal of Image and Graphics, 2020, 25 (10): 2259- 2270

DOI:10.11834/jig.200187      [本文引用: 1]

NI Z L, BIAN G B, ZHOU X H, et al. RAUNet: residual attention U-Net for semantic segmentation of cataract surgical instruments [C]// International Conference on Neural Information Processing. Shenzhen: Springer, 2019: 139-149.

[本文引用: 1]

OKTAY O, SCHLEMPER J, FOLGOC L L, et al. Attention U-Net: learning where to look for the pancreas [EB/OL]. [2021-08-01]. https://arxiv.org/abs/1804.03999.

[本文引用: 1]

PARK J, WOO S, LEE J Y, et al. BAM: bottleneck attention module [EB/OL]. [2021-08-01]. https://arxiv.org/abs/1807.06514.

[本文引用: 1]

WOO S, PARK J, LEE J Y, et al. CBAM: convolutional block attention module [C]// Proceedings of the European Conference on Computer Vision. Munich: Springer, 2018: 3-19.

[本文引用: 2]

GU R, WANG G, SONG T, et al

CA-Net: comprehensive attention convolutional neural networks for explainable medical image segmentation

[J]. IEEE Transactions on Medical Imaging, 2020, 40 (2): 699- 711

[本文引用: 2]

DOSOVITSKIY A, BEYER L, KOLESNIKOV A, et al. An image is worth 16×16 words: Transformers for image recognition at scale [EB/OL]. [2021-08-01]. https://arxiv.org/abs/2010.11929.

[本文引用: 1]

LIU Z, LIN Y, CAO Y, et al. Swin Transformer: hierarchical vision transformer using shifted windows [C]// Proceedings of the IEEE/CVF International Conference on Computer Vision. Montreal: IEEE, 2021: 10012-10022.

[本文引用: 2]

GUO J, HAN K, WU H, et al. CMT: convolutional neural networks meet vision transformers [EB/OL]. [2021-08-01].https://arxiv.org/abs/2107.06263.

[本文引用: 2]

CAO H, WANG Y, CHEN J, et al. Swin-Unet: Unet-like pure Transformer for medical image segmentation [EB/OL]. [2021-08-01]. https://arxiv.org/abs/2105.05537.

[本文引用: 1]

GAO Y, ZHOU M, METAXAS D. UTNet: a hybrid Transformer architecture for medical image segmentation [EB/OL]. [2021-08-01]. https://arxiv.org/abs/2107.00781.

[本文引用: 1]

CODELLA N C F, GUTMAN D, CELEBI M E, et al. Skin lesion analysis toward melanoma detection: a challenge at the 2017 international symposium on biomedical imaging (ISBI), hosted by the international skin imaging collaboration (ISIC) [C]// 2018 IEEE 15th international symposium on biomedical imaging (ISBI 2018). Washington DC: IEEE, 2018: 168-172.

[本文引用: 1]

BERNAL J, SANCHEZ J, VILARINO F

Towards automatic polyp detection with a polyp appearance model

[J]. Pattern Recognition, 2012, 45 (9): 3166- 3182

DOI:10.1016/j.patcog.2012.03.002      [本文引用: 1]

HE K, ZHANG X, REN S, et al. Deep residual learning for image recognition [C]// Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. Las Vegas: IEEE, 2016: 770-778.

[本文引用: 1]

WU H, PAN J, LI Z, et al

Automated skin lesion segmentation via an adaptive dual attention module

[J]. IEEE Transactions on Medical Imaging, 2020, 40 (1): 357- 370

[本文引用: 1]

SINHA A, DOLZ J

Multi-scale self-guided attention for medical image segmentation

[J]. IEEE Journals on Biomedical and Health Informatics, 2021, 25 (1): 121- 130

DOI:10.1109/JBHI.2020.2986926      [本文引用: 1]

/