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^l_i, O^l_i) (Iil,Oil)来表示映射,其中 I i l ∈ I I^l_i∈I Iil∈I表示输入张量, O i l ∈ O O^l_i∈O Oil∈O表示候选操作。输入张量集包括前一个单元的输出 H l − 1 H^{l−1} Hl−1,前一个单元的输出 H l − 2 H^{l−2} Hl−2,以及当前单元中前块的输出, H 1 l , … , H i − 1 l H^l_1,…,H^l_{i−1} H1l,…,Hi−1l。候选操作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)
其中
P
a
r
t
j
→
i
Part_{j→i}
Partj→i为通道选择的采样掩码,
α
α
α为归一化标量,称为体系结构参数,表示确定候选操作的权重。为了保证反向传播的可微性,实现了softmax。最后,输出张量
H
l
=
C
o
n
c
a
t
(
H
i
l
∣
i
<
=
B
)
H^l = Concat({H^l_i|i <= B})
Hl=Concat(Hil∣i<=B),其中Concat(·)是连接,B是块的数量。Cell级搜索表示为
H
l
=
C
e
l
l
(
H
l
−
1
,
H
l
−
2
;
α
)
H^l = Cell(H^{l−1},H^{l−2};α)
Hl=Cell(Hl−1,Hl−2;α)。
出层(Out Layer):层级搜索的目的是针对特定的数据集,在网络搜索空间内寻找最优的网络骨干网,具体步骤是将不同分辨率的特征结合起来,更好地提取特征。在过去的研究中[15,20,21],常规算法是使用另一个结构参数β进行线性组合特征融合。
然而, β β β的量并不大,通常在数百个,但负责复杂的特征融合,特别是在应用部分通道连接后。通过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(2sHl−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
∈
R
H
×
W
×
C
x∈\mathbb{R}^{H×W ×C}
x∈RH×W×C,其中
H
×
W
H×W
H×W表示空间分辨率,C表示通道数。首先,执行标记化[22]。将
x
x
x重塑为
x
p
i
∈
R
P
×
P
×
C
∣
i
=
1
,
…
N
{x^i_p∈R^{P ×P ×C}|i = 1,…N}
xpi∈RP×P×C∣i=1,…N通过
P
×
P
P × P
P×P卷积(stride = P),其中
N
=
H
W
/
P
2
N = HW /P^2
N=HW/P2。二是patch嵌入和Transformer。将patch
x
p
x_p
xp映射到d维嵌入空间中,并添加特定位置嵌入以保持位置信息,然后应用transformer,包括多头自注意(MSA)和多层感知器(MLP)块:
z
l
∗
=
M
S
A
(
L
N
(
z
l
−
1
)
)
+
z
l
−
1
z
l
=
M
L
P
(
L
N
(
z
l
∗
)
)
+
z
l
∗
\begin{align} \mathbf{z}_l^*&=MSA(LN(\mathbf{z}_{l-1}))+\mathbf{z}_{l-1}\tag{3} \\ \mathbf{z}_l&=MLP(LN(\mathbf{z}_l^*))+\mathbf{z}_l^*\tag{4} \end{align}
zl∗zl=MSA(LN(zl−1))+zl−1=MLP(LN(zl∗))+zl∗(3)(4)
式中LN(·)表示层归一化。
因此,Transformer输出为 z l ∈ R H W P 2 × D \mathbf{z}_l\in\mathbb{R}^{\frac{HW}{P^2}\times D} zl∈RP2HW×D,然后我们将其重塑为 R H P × W P × D \mathbb{R}^{\frac{H}{P}\times\frac{W}{P}\times D} RPH×PW×D用于解码器。在这里,我们需要将从Transformers中提取的全局特征与从NAS中提取的局部特征结合起来。为了避免丢失低级或高级细节,使用U形解码器是必要的。
我们首先使用2D卷积来调整来自Transformer的编码特征的通道,使其具有与用于特征聚合的 s H l ^sH^l sHl相同数量的通道。由于我们不期望Transformer带来过多的参数,因此H/P和W/P很小,因此应将其上采样到s分辨率的空间尺寸。在对图像进行整形和重采样后,将两部分特征进行拼接,然后用卷积层将下一分辨率的通道数匹配为 s 2 H l \frac{s}{2}H^{l} 2sHl通道数,从而得到当前分辨率的输出。
同样,对下一个分辨率特征再次执行“Upsample-Concatenation-Convolution”(如图1©所示),直到合并s = 4个分辨率特征。为了分割目的,最后使用上样本层和卷积层将特征恢复到全分辨率和特定的类数,以预测密集输出。
2.3. Semi-Supervised学习
由于医学图像数据集总是有少量的标记图像和大量的未标记数据。因此,我们尝试将NAS与Mean Teacher方法相结合[23],以增强模型的有效性和泛化性。
对于Mean Teacher,我们有两个模型:
S
t
u
d
e
n
t
f
(
θ
s
)
Student \ f(θ_s)
Student f(θs)和
T
e
a
c
h
e
r
f
(
θ
t
)
Teacher \ f(θ_t)
Teacher f(θt),其中
θ
s
θ_s
θs和
θ
t
θ_t
θt分别表示Student和Teacher的网络参数。
θ
t
θ_t
θt由
θ
s
θ_s
θs的指数移动平均线(EMA)更新:
θ
t
,
i
=
α
θ
s
,
i
−
1
+
(
1
−
α
)
θ
s
,
i
(5)
\theta_{t,i}=\alpha\theta_{s,i-1}+(1-\alpha)\theta_{s,i} \tag{5}
θt,i=αθs,i−1+(1−α)θs,i(5)
其中
α
α
α是控制更新速度的超参数,
I
I
I是迭代次数。给定标记数据集
D
l
D^l
Dl和未标记数据集
D
u
D^u
Du,定义一致性正则化来更新
θ
s
θ_s
θs:
L
c
=
∑
D
l
l
M
S
E
(
S
(
p
l
t
)
,
S
(
p
l
s
)
)
m
+
∑
D
u
l
M
S
E
(
S
(
p
u
t
)
,
S
(
p
u
s
)
)
n
(6)
\mathcal{L}_{c}=\frac{\sum_{{\cal D}_{l}}l_{MSE}(S(p_{l}^{t}),S(p_{l}^{s}))}{m}+\frac{\sum_{{\cal D}_{u}}l_{MSE}(S(p_{u}^{t}),S(p_{u}^{s}))}{n} \tag{6}
Lc=m∑DllMSE(S(plt),S(pls))+n∑DulMSE(S(put),S(pus))(6)
其中
p
l
t
p^t_l
plt和
p
l
s
p^s_l
pls为老师和学生在标记图像上的输出,
p
u
t
p^t_u
put和
p
u
s
p^s_u
pus为未标记图像上的输出,
l
M
S
E
l_{M SE}
lMSE为均方误差,S(·)为用于通道维度的softmax函数,用于尺度控制。
Student的监督损失由Teacher的输出和标记图像的真值计算:
L
s
=
1
m
∑
D
l
l
C
E
(
p
l
t
,
y
l
)
(7)
\mathcal{L}_s=\frac{1}{m}\sum_{\mathcal{D}_l}l_{CE}(p_l^t,y_l) \tag{7}
Ls=m1Dl∑lCE(plt,yl)(7)
其中
l
C
E
l_{CE}
lCE为交叉熵损失函数。因此,用于训练
θ
s
θ_s
θs的总损失是一致性正则化和监督损失的总和:
L
t
o
t
a
l
=
λ
0
L
s
+
λ
1
L
c
L_{total} = λ_0L_s+λ_1L_c
Ltotal=λ0Ls+λ1Lc,其中
λ
0
=
1
λ_0 = 1
λ0=1,
λ
1
λ_1
λ1遵循[24]中的指数上升函数。
2.4. Optimization
为了构造一个可微的计算图,我们对结构参数 α α α使用连续松弛来控制内cell的连接,对结构参数 γ γ γ使用微分卷积层来控制外层的连接,这使得梯度下降成为可能。在训练中,我们将标记和未标记的训练数据集分别分成两个不相交的集: D l , A D_{l,A} Dl,A和 D l , B D_{l,B} Dl,B, D u , A D_{u,A} Du,A和 D u , B D_{u,B} Du,B。每个epoch的优化可以总结为:
- 通过 ∇ w L D l , A , D u , A ( w , α , γ ) \nabla_w \mathcal{L}_{\mathcal{D}_{l, A}, \mathcal{D}_{u, A}}(w, \alpha, \gamma) ∇wLDl,A,Du,A(w,α,γ)使用 D l , A \mathcal{D}_{l, A} Dl,A和 D u , A \mathcal{D}_{u, A} Du,A更新 f ( θ s ) f\left(\theta_s\right) f(θs)的权重参数 w w w
- 通过EMA 更新 f ( θ t ) f\left(\theta_t\right) f(θt)的权重参数 w w w
- 通过 ∇ α , γ L D l , B , D u , B ( w , α , γ ) \nabla_{\alpha, \gamma} \mathcal{L}_{\mathcal{D}_{l, B}, \mathcal{D}_{u, B}}(w, \alpha, \gamma) ∇α,γLDl,B,Du,B(w,α,γ)使用 D l , B \mathcal{D}_{l, B} Dl,B 和 D u , B \mathcal{D}_{u, B} Du,B更新f(θs)的结构参数 α \alpha α a和 γ \gamma γ
- 通过EMA更新 f ( θ t ) f\left(\theta_t\right) f(θt)的体系结构参数 α \alpha α 和 γ \gamma γ
其中 L \mathcal{L} 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实现了比现有方法更优越的性能。