SSHNN: SEMI-SUPERVISED HYBRID NAS NETWORK FOR ECHOCARDIOGRAPHIC IMAGE SEGMENTATION

SSHNN: 超声心动图图像分割的半监督混合NAS网络

论文链接：https://arxiv.org/abs/2309.04672

项目链接：Code will be made publicly available

ABSTRACT

准确的医学图像分割，特别是对于含有不可忽略噪声的超声心动图图像，需要精心设计网络。与手工设计相比，Neural Architecture Search (NAS)由于搜索空间更大、自动优化，实现了更好的分割效果，但现有的大多数方法在分层特征聚合方面较弱，采用“强编码器、弱解码器”的结构，不足以处理全局关系和局部细节。为了解决这些问题，我们提出了一种新的半监督混合NAS网络，用于精确的医学图像分割，称为SSHNN。在SSHNN中，我们创造性地在分层特征融合中使用卷积运算，而不是标准化的标量，以避免丢失细节，使NAS成为更强大的编码器。引入Transformer对全局上下文进行补偿，设计u型解码器，将全局上下文与局部特征有效连接。具体来说，我们实现了一种半监督算法Mean-Teacher来克服标记医学图像数据集的有限体积问题。在CAMUS超声心动图数据集上的大量实验表明，SSHNN优于最先进的方法，实现了准确的分割。

1. INTRODUCTION

二维超声心动图是临床常规测量心脏形态和功能，进而获得诊断的重要医学成像技术[1]。由于光照不均匀引起的噪声、伪影和低对比度，超声心动图图像中的器官边界模糊不清[2]。此外，超声心动图图像中的结构在不同尺度下呈现出不同的特征，需要该模型有效处理医学图像分割的多尺度特征。卷积神经网络(Convolutional Neural Networks, CNN)得益于感受野，具有非凡的特征提取能力，被广泛采用[3-5]。特别是在医疗领域，U型编码器-解码器网络的提出推动了语义分割的发展[6-9]。此外，将Transformer[10]引入计算机视觉并转换为视觉Transformer (ViT)[11]以增强全局接受能力时，重点关注U型网络与Transformer的结合，如TransUNet[12]、Swin Transformer[13]等。人工设计的网络在提高分割精度的同时也变得更加复杂。

为了解决网络设计的难度和复杂性，提出了神经架构搜索(NAS)，旨在实现网络设计的自动化和精确化。其中，NAS-UNet[14]用可搜索单元替换CNN部分，减少网络参数。为了扩大搜索空间更好的泛化，HNAS (Hierarchical NAS)也在层级进行搜索，如Auto-DeepLab [15]， Dynamic routing[16]等。此外，HCT-Net[17]采用ViT添加上下文信息。然而，这些方法在分层优化中采用归一化标量进行特征融合，容易因参数不足而导致局部细节丢失。此外，解码器虽然学习了可用的多尺度特征，但只依赖于拼接层和卷积层，导致“强编码器，弱解码器”，无法捕获足够的分割信息。因此，我们的重点是通过范例手工设计和HNAS的结合来增强模型分割能力。

在本文中，我们考虑到标记医学图像数据集的大小限制，提出了一种具有半监督学习的新型混合NAS，称为SSHNN。在HNAS的设计中，我们用卷积代替简单的标量进行局部特征聚合，追求更高程度的灵活性。为了克服“强编码器，弱解码器”的问题，在HNAS后增加U形解码器，将多尺度特征恢复到原始尺寸，并利用Transformer对全局上下文进行补偿。实验证明，在超声心动图公共医学图像数据集CAMUS上，SSHNN优于最先进的SOTA方法[1]。

2. METHODOLOGIES

在本节中，我们从HNAS设计、解码器、半监督学习和优化策略四个部分来阐述SSHNN。网络框架如图1所示。

2.1. HNAS 设计

Inner Cell(内cell)： cell搜索空间可以用一个有向无环图来表示，由块和边组成，分别表示从输入张量到1个输出张量的映射，以及候选操作。对于第l层单元格的第I块，定义一个元组$(I_i^l,O_i^l)$来表示映射，其中$I_i^l\in I$表示输入张量，$O_i^l\in O$表示候选操作。输入张量集包括前一个单元的输出$H^{l-1}$，前一个单元的输出$H^{l-2}$，以及当前单元中前块的输出，$H_1^l,\dots ,H_{i-1}^l$。候选操作O集包括：深度可分卷积、空洞卷积、平均池化、最大池化、跳过连接和无连接。

为了降低存储成本，采用部分通道连接[18]，将1/n部分维度特征发送到cell，其余特征保持不变。此外，在追求可微搜索空间时重用连续松弛[19]，然后可以应用随机梯度下降法(SGD)。则block的输出张量为：
H i l = ∑ H j l ∈ I O k l ∈ O e x p { α j → i k } ∑ m = 1 ∣ O ∣ e x p { α j → i m } ⋅ O k l ( P a r t j → i ∘ H j l ) + ( 1 − P a r t j → i ) ∘ H i l (1) \begin{aligned} H_{i}^{l}& =\sum_{H_{j}^{l}\in\mathcal{I}O_{k}^{l}\in\mathcal{O}}\frac{exp\{\alpha_{j\rightarrow i}^{k}\}}{\sum_{m=1}^{|\mathcal{O}|}exp\{\alpha_{j\rightarrow i}^{m}\}}\cdot O_{k}^{l}(Part_{j\rightarrow i}\circ H_{j}^{l}) \\ &+(1-Part_{j\rightarrow i})\circ H_{i}^{l} \end{aligned} \tag{1} Hil​​=Hjl​∈IOkl​∈O∑​∑m=1∣O∣​exp{αj→im​}exp{αj→ik​}​⋅Okl​(Partj→i​∘Hjl​)+(1−Partj→i​)∘Hil​​(1)
其中$Par{t}_{j\to i}$为通道选择的采样掩码，$\alpha$为归一化标量，称为体系结构参数，表示确定候选操作的权重。为了保证反向传播的可微性，实现了softmax。最后，输出张量$H^l=Concat(H_i^l|i<=B)$，其中Concat(·)是连接，B是块的数量。Cell级搜索表示为$H^l=Cell(H^{l-1},H^{l-2};\alpha )$。

出层(Out Layer)：层级搜索的目的是针对特定的数据集，在网络搜索空间内寻找最优的网络骨干网，具体步骤是将不同分辨率的特征结合起来，更好地提取特征。在过去的研究中[15,20,21]，常规算法是使用另一个结构参数β进行线性组合特征融合。

然而，$\beta$的量并不大，通常在数百个，但负责复杂的特征融合，特别是在应用部分通道连接后。通过3.4节的估计，线性组合不能完全融合多尺度特征，因为数百个通道代表不同的特征，需要更多的参数来阐述。

因此，我们将归一化标量替换为卷积核，以实现更高程度的灵活性，如图1(d，左)所示。此外，由于在图像分割方面具有令人难以置信的能力，因此利用了长和短跳跃连接[7,21]。

层设计的整体表示：
s H l = C o n v ( C o n c a t ( D o w n ( C e l l ( s 2 H l − 1 , s H l − 2 ; α ) ) , U p ( C e l l ( 2 s H l − 1 , s H l − 2 ; α ) ) , C e l l ( s H l − 1 , s H l − 2 ; α ) , { s H l ′ ∈ s H ∣ l ′ < l } ) ) (2) \begin{aligned} ^{s}H^{l}& =Conv(Concat(Down(Cell(^{\frac{s}{2}}H^{l-1},^{s}H^{l-2};\alpha)), \\ &Up(Cell(^{2s}H^{l-1},^{s}H^{l-2};\alpha)),Cell(^{s}H^{l-1},^{s}H^{l-2};\alpha), \\ &\{^{s}H^{l^{\prime}}\in^{s}H|l^{\prime}<l\})) \end{aligned} \tag{2} sHl​=Conv(Concat(Down(Cell(2s​Hl−1,sHl−2;α)),Up(Cell(2sHl−1,sHl−2;α)),Cell(sHl−1,sHl−2;α),{sHl′∈sH∣l′<l}))​(2)
式中，Conv(·)表示卷积，将融合特征变换为s分辨率的相同通道数，Down(·)表示下采样，U p(·)表示上采样。

2.2. Decoder

我们不是简单地连接多尺度特征，然后通过卷积层进行处理，在搜索最佳网络结构后将特征恢复到原始图像大小，其中这些层充当解码器，HNAS本质上充当编码器，而是应用ViT添加全局上下文并使用U形解码器结构，如图1(a)所示。

定义输入图像$x\in {\mathbb{R}}^{H\times W\times C}$，其中$H\times W$表示空间分辨率，C表示通道数。首先，执行标记化[22]。将$x$重塑为$x_p^i\in R^{P\times P\times C}|i=1，\dots N$通过$P\times P$卷积(stride = P)，其中$N=HW/P^2$。二是patch嵌入和Transformer。将patch $x_p$映射到d维嵌入空间中，并添加特定位置嵌入以保持位置信息，然后应用transformer，包括多头自注意(MSA)和多层感知器(MLP)块：
$$\begin{array}{rl}{\mathbf{z}}_l^{\ast }& =MSA(LN({\mathbf{z}}_{l-1}))+{\mathbf{z}}_{l-1}\\ {\mathbf{z}}_l& =MLP(LN({\mathbf{z}}_l^{\ast }))+{\mathbf{z}}_l^{\ast }\end{array}\text{\hspace{1em}(3)(4)}$$
式中LN(·)表示层归一化。

因此，Transformer输出为${\mathbf{z}}_l\in {\mathbb{R}}^{\frac{HW}{P^2}\times D}$，然后我们将其重塑为${\mathbb{R}}^{\frac{H}{P}\times \frac{W}{P}\times D}$用于解码器。在这里，我们需要将从Transformers中提取的全局特征与从NAS中提取的局部特征结合起来。为了避免丢失低级或高级细节，使用U形解码器是必要的。

我们首先使用2D卷积来调整来自Transformer的编码特征的通道，使其具有与用于特征聚合的${}^s{H}^l$相同数量的通道。由于我们不期望Transformer带来过多的参数，因此H/P和W/P很小，因此应将其上采样到s分辨率的空间尺寸。在对图像进行整形和重采样后，将两部分特征进行拼接，然后用卷积层将下一分辨率的通道数匹配为$\frac{s}{2}{H}^l$通道数，从而得到当前分辨率的输出。

同样，对下一个分辨率特征再次执行“Upsample-Concatenation-Convolution”(如图1©所示)，直到合并s = 4个分辨率特征。为了分割目的，最后使用上样本层和卷积层将特征恢复到全分辨率和特定的类数，以预测密集输出。

2.3. Semi-Supervised学习

由于医学图像数据集总是有少量的标记图像和大量的未标记数据。因此，我们尝试将NAS与Mean Teacher方法相结合[23]，以增强模型的有效性和泛化性。

对于Mean Teacher，我们有两个模型：$Studentf({\theta }_s)$和$Teacherf({\theta }_t)$，其中${\theta }_s$和${\theta }_t$分别表示Student和Teacher的网络参数。${\theta }_t$由${\theta }_s$的指数移动平均线(EMA)更新：
$${\theta }_{t,i}=\alpha {\theta }_{s,i-1}+(1-\alpha ){\theta }_{s,i}\text{\hspace{1em}(5)}$$
其中$\alpha$是控制更新速度的超参数，$I$是迭代次数。给定标记数据集$D^l$和未标记数据集$D^u$，定义一致性正则化来更新${\theta }_s$：
$${\mathcal{L}}_c=\frac{{\sum }_{{\mathcal{D}}_l}{l}_{MSE}(S(p_l^t),S(p_l^s))}{m}+\frac{{\sum }_{{\mathcal{D}}_u}{l}_{MSE}(S(p_u^t),S(p_u^s))}{n}\text{\hspace{1em}(6)}$$
其中$p_l^t$和$p_l^s$为老师和学生在标记图像上的输出，$p_u^t$和$p_u^s$为未标记图像上的输出，$l_{MSE}$为均方误差，S(·)为用于通道维度的softmax函数，用于尺度控制。

Student的监督损失由Teacher的输出和标记图像的真值计算：
$${\mathcal{L}}_s=\frac{1}{m}\sum _{{\mathcal{D}}_l}{l}_{CE}(p_l^t,y_l)\text{\hspace{1em}(7)}$$
其中$l_{CE}$为交叉熵损失函数。因此，用于训练${\theta }_s$的总损失是一致性正则化和监督损失的总和：$L_{total}={\lambda }_0{L}_s+{\lambda }_1{L}_c$，其中${\lambda }_0=1$， ${\lambda }_1$遵循[24]中的指数上升函数。

2.4. Optimization

为了构造一个可微的计算图，我们对结构参数$\alpha$使用连续松弛来控制内cell的连接，对结构参数$\gamma$使用微分卷积层来控制外层的连接，这使得梯度下降成为可能。在训练中，我们将标记和未标记的训练数据集分别分成两个不相交的集：$D_{l,A}$和$D_{l,B}$, $D_{u,A}$和$D_{u,B}$。每个epoch的优化可以总结为：

1. 通过${\nabla }_w{\mathcal{L}}_{{\mathcal{D}}_{l,A},{\mathcal{D}}_{u,A}}(w,\alpha ,\gamma )$使用${\mathcal{D}}_{l,A}$和${\mathcal{D}}_{u,A}$更新$f\left({\theta }_s\right)$的权重参数$w$
2. 通过EMA 更新$f\left({\theta }_t\right)$的权重参数$w$
3. 通过${\nabla }_{\alpha ,\gamma }{\mathcal{L}}_{{\mathcal{D}}_{l,B},{\mathcal{D}}_{u,B}}(w,\alpha ,\gamma )$使用${\mathcal{D}}_{l,B}$ 和 ${\mathcal{D}}_{u,B}$更新f(θs)的结构参数$\alpha$ a和 $\gamma$
4. 通过EMA更新$f\left({\theta }_t\right)$的体系结构参数$\alpha$ 和 $\gamma$

其中$\mathcal{L}$是分割小批的总损失。

3. EXPERIMENTAL RESULTS

3.1. 数据集和评估指标

CAMUS[1]数据集用于评估SSHNN的性能，CAMUS数据集是二维超声心动图中的开放大规模数据集，收集自500例患者。注意，从每个病人身上收集了一个四室和一个两室的视图序列，除了舒张末期(ED)和收缩末期(ES)的时刻外，持续了大约20张照片，没有标签。因此，可访问的标记数据集有2000张用于监督损失计算的超声心动图图像，未标记数据集有大约19000张超声心动图图像。分割标签有四种类型，分别由手工标注的左心室心内膜，心肌，左心房和背景。本文采用交联(IoU)、骰子系数(Dice)和参数(Params)作为评价指标。

3.2. 实现细节

我们使用大小为256 × 256的图像作为网络输入。每个单元有B = 5个块。使用滤波器乘法器F和层数L控制网络复杂度，F初始化为8。因此，在s = 4中，有B ×F × s 4 = 40个过滤器。当空间大小从s减小到2s时，过滤器的数量翻倍。我们使用4个transformer来提取全局上下文。实验在单个Nvidia RTX 3090Ti GPU上进行。对于权值w和架构γ的训练，应用的优化器是SGD，动量为0.9，权值衰减为0.0003，初始值为0.01。对于架构α， Adam[27]的学习率为0.003，权重衰减为0.001。总epoch数为40，由于w的不稳定更新使其陷入局部最优，架构在10个epoch后进行优化。教师与学生同步更新，稳定训练。最后，使用10%的标记图像进行验证，其余的使用未标记的图像进行训练。

3.3. 实验结果

为了证明SSHNN的有效性，我们将其与SOTA方法进行了比较。表1和图2分别阐述了每个网络在CAMUS数据集上的数值和视觉结果。每个指标的均值和标准差值是通过对数据集的10倍交叉验证获得的。从SSHNN系列来看，显然L越大，模型容量越大，分割性能越好，但代价是参数越多(速度越慢)。与表1中Dice次优的输血相比，SSHNN-L的Dice提高了0.98%，IoU增强了1.05%。同样的，SSHNN-M也具有用更少的参数做出更好的Dice，验证了其优越性。

3.4. 消融实验

如表2所示，我们进行了多组测试来评估卷积融合、U形解码器和Transformer以及标记图像的比例的影响。

首先，当未标记图像与标记图像Nu/Nl = 10时，经过卷积融合的SSHNN-L的Dice比没有经过卷积融合的SSHNN-L高1.00%。此外，使用U型解码器和Transformer也使Dice在这个分数不变的情况下增加了1.33%，验证了它们的有效性。显然，将卷积和U形解码器与Transformer结合使用效果更好，当Nu/Nl = 10时，SSHNN-L仍为0.921 Dice，符合以最小监督进行学习的半监督目标。请注意，当Nu/Nl = 2时，该模型的性能优于其他模型，因为未标记的数据提供了额外的信息，并且不会掩盖标记数据的特征。SSHNN-S和SSHNN-M的结果再次证明了该方法的鲁棒性。

最后但并非最不重要的是，表3讨论了F的影响，表明更高的模型容量以昂贵的内存成本促进分割能力，这是一个权衡问题。

4. CONCLUSION

本文提出了一种用于医学图像分割的半监督混合NAS网络SSHNN。采用卷积融合的方法，充分利用HNAS分层搜索提取的多尺度特征。此外，解码部分还引入了U型解码器和Transformer。实现了半监督学习方法Mean Teacher，克服了标注数据量有限的问题。最后，在CAMUS数据集上的实验表明，SSHNN实现了比现有方法更优越的性能。