HF-UNet
关键词: 医学图像, 多任务学习, 注意力机制
进度: Yes
一、来源
IEEE Transactions on Medical Imaging
Early Access 13 April 2021
原文链接：https://arxiv.org/abs/2005.10439
二、标题
HF-UNet： Learning Hierarchically Inter-Task Relevance in Multi-Task U-Net for Accurate Prostate Segmentation in CT images
三、作者及单位
Kelei He, Chunfeng Lian, Bing Zhang, Xin Zhang, Xiaohuan Cao, Dong Nie, Yang Gao, Junfeng Zhang, Dinggang Shen
南京大学
四、方法简介
本文采用两阶段网络技术解决了CT图像前列腺分割。第一阶段是快速定位，第二阶段是精确分割。在第二阶段，为了精确的分割前列腺，本文将前列腺分割制定为一个多任务学习框架，主要任务是分割前列腺，辅助任务是勾画前列腺边界。在这里，辅助任务是为CT图像中前列腺边界不清提供额外的指导。
传统的多任务深度网络通常在所有任务之间共享大部分参数(即特征表示)，这可能会限制其数据拟合能力，因为不可避免地忽略了不同任务的特殊性。与此相反，本文提出采用层次融合的U-Net结构，即HFUNet来解决这些问题。HF-UNet有两个互补的分支用于两个任务，并提出了基于注意力的任务一致性学习块（TCL）来实现两个编解码解码分支之间的信息交换。因此，HF-UNet赋予了分层学习不同任务共享表征的能力，同时保持了学习到的不同任务表征的特异性。
定位
本文首先采用区域定位网络对包含整个前列腺区域的子图像进行裁剪，然后通过一个2D U-Net架构的网络自动定位前列腺区域。
根据获得的器官中心裁剪128 × 128 × 128的区域，以完全覆盖前列腺器官。然后，从识别区域裁剪出的patch用于训练后续的网络HF-UNet，以实现更准确的前列腺分割和更高效的网络训练。
分割
本文呢提出的HF-UNet的架构如下图所示。为了实现前列腺的精确分割，HF-UNet设置了两个任务:1)前列腺的分割， 2)前列腺边界的勾画。这里引入轮廓感知子网络任务，为该网络提供重要的形态学指导，有助于更好地区分前列腺边界。由于传统fnc中使用的降采样操作，通常会破坏前列腺边界体素之间的空间相关性，导致最终的分割不能描绘完整的前列腺轮廓。
Figure 1. The pipeline of the proposed segmentation framework.
Figure 2. 任务一致性学习模块（TCL）
利用提出的TCL块对两个分支的信息进行组合学习。在这里，TCL块可以看作是从全局的角度(即任务级)对两个分支学习到的知识进行建模。每个分支中学习到的特征可以看作是从局部角度(即样本层次)学习到的知识。本文构建了两种TCL块:
-
加权残差TCL块，用于评价共享信息和任务特定信息的有效性，
-
基于注意的TCL，自适应学习上述两种信息的比例。
损失函数
通过多任务学习，共同解决了三个任务：分割、回归和任务一致性学习。
L
=
λ
1
L
c
l
s
+
λ
2
L
r
e
g
+
λ
3
L
t
c
l
L=\lambda_{1}L_{cls}+\lambda_{2}L_{reg}+\lambda_{3}L_{tcl}
L=λ1Lcls+λ2Lreg+λ3Ltcl
回归任务往往会产生巨大的损失。因此，本文使用两个技巧来解决这个问题。
（1）在训练阶段开始时，只训练分割分支的参数，固定回归分支和TCL块中的所有参数。这样可以更好地初始化多任务网络的权值。
（2）固定总损失函数中的权值： λ 1 \lambda_{1} λ1= 1.0， λ 2 \lambda_{2} λ2= 0.01， λ 3 \lambda_{3} λ3= 1.0
五、实验数据
在一个包含339例患者的CT图像数据集上评估了所提出的方法。数据集包含339张图像，图像大小为512 × 512 ×(61 ~ 508)。这些图像的平面分辨率为0.932 ~ 1.365 mm，切片厚度为1 ~ 3mm。
六、主要结果
实验结果表明，HF-UNet的性能优于传统的多任务网络体系结构和最新的多任务网络方法。
这表明该方法可以生成更稳健的前列腺掩膜形状，更接近实际的器官表面。这主要是由于形态学表示的结合，这有助于更好地描绘器官边界。同时，所提出的分层融合策略可以从网络的低层次特征中获取详细的轮廓信息