让UNet火力全开 | All Attention出击UNet让ResNet50+UNet提升31.8%

点击下方卡片，关注「AI视界引擎」公众号

脑出血在 Head CT扫描中作为第一线工具，帮助专家诊断不同类型的出血。然而，在同一类型中，它们的形状各异，但在形状、大小和位置上相似且令人困惑。

为了解决这个问题，本文提出了一种全注意力U-Net。它在U-Net编码器侧使用通道注意力来增强特定类别的特征提取，而在U-Net解码器侧使用空间和通道注意力进行更准确的形状提取和类型分类。

模拟结果显示，与 Baseline 、ResNet50+U-Net相比，提高了31.8%，并且在有限注意力情况下表现更好。

I Introduction

Head CT扫描是急诊室广泛接受的第一线工具，用于处理 Head 损伤、中风或其他颅内病变，因为它具有短的获取时间。识别个体病变类型和区域对于诊断非常重要，但也非常耗时，尤其是对于小而紧急的疾病。近年来，利用深度学习进行医学图像分析是一种有前景且表现良好的方法，已经在识别糖尿病性视网膜病变、分类皮肤病变、胸部CT异常检测、X射线图像和 Head CT等方面得到应用。

对于作者的目标问题， Head CT扫描中的颅内出血，语义分割可以帮助定位和量化这些疾病模式。然而，由于 Head CT扫描中的颅内出血有七种类型，且同一种类型的出血形状不同，不同类型的出血形状相似，大小和位置相似，这给深度学习网络带来了挑战。

Head CT扫描中的颅内出血的语义分割也面临着高分辨率输入（高达512×512）的挑战，由于某些极小的病变尺寸，无法通过调整图像大小来进行训练和测试。此外，这也阻止了使用某些需要大量内存进行训练的高级模型，从而使得批量大小变得太小甚至无法训练。在深度学习模型中，U-Net已经广泛应用于医学图像分析，由于其较小的训练成本和即使在少量数据情况下表现出色，因此在医学图像案例中得到了广泛应用。然而，原始U-Net架构由堆叠卷积组成，没有增强不同通道和层信息的依赖性。这使得很难提取全局特征并生成正确的形状进行医学图像分析。

为了提高性能，已经提出了各种增强版本的U-Net。一种方法是采用最新的分类网络作为U-Net架构的 Backbone [11][12][13]。另一种方法是使用不同的聚合方法将不同 Level 的特征结合。然而，这些聚合方法，其中聚合单元通过一系列嵌套和密集的跳过路径连接，消耗过多的计算和内存。

此外，对于目标问题，它们的表现仍然受到限制，如作者的实验结果所示。另一种方法是使用可以增加感受野并改善模型中长程依赖关系的注意力机制，而无需添加太多计算和参数。这些注意力网络包括在空间、通道或组维度上的注意力，已集成到U-Net解码器侧。然而，解码器侧输入特征图的通道映射仍然被破坏，这需要进一步改进，尤其是对于变异率高的病变数据集。此外，没有人考虑了编码器侧的注意力问题。

为了解决上述问题，作者提出了全注意力U-Net，通过在U-Net编码器侧使用通道注意力和在U-Net解码器侧使用空间和通道注意力来解决形状、大小和位置之间的冲突。模拟结果显示，所提出的网络可以预测各种颅内出血，但仍具有较小的模型大小，这优于 Baseline ResNet50+U-Net和仅有一种注意力U-Net。

II Dataset

该数据集包含总共51个案例，每个案例约2,048张来自台湾长庚纪念医院的 Head CT扫描图像，这些CT扫描图像均由专业医生进行标记。数据使用得到了当地研究伦理委员会的批准。切片数量在每例案例中从30到50不等。这些案例包括7种颅内出血类型，即脑内出血（ICH）、急性硬膜下血肿（SDH）、蛛网膜下腔出血（SAH）、外膜下血肿（EDH）、慢性硬膜下血肿（CSDH）、气颅、脑室内出血（IVH）。

Data Imbalance And High Resolution Images

与其他医学图像数据集一样，由于 GT 标记的困难，作者的数据集也相对较小。此外，一个案例中的许多扫描显示没有病变（约占46.55%），如图1所示，病变区域在大多数扫描中都很小。

这种数据不平衡会导致训练困难。此外，CT扫描图像的大小较大，为512 x 512。高分辨率输入将消耗过多的GPU内存并延长训练时间。然而，由于某些类型的病变区域（如图2所示，不到20像素）在缩放图像时将被移除，因此无法重新缩放这些图像。

High Variety of Lesions

图1显示了不同颅内出血的示例。一张图像可能包含一到几种病变。但是，如果同时出现类似的病变，将很难区分，即使是专家。

例如，EDH和SDH在质地和位置上相似（如图3所示），唯一的区别是它们的形状。其他令人困惑的病变有SAH、IVH和ICH，在某些情况下形状和大小相似。

此外，即使在同一类型的病变中，其形状和大小也有很多变化，如图1所示的ICH。CSDH可以通过其纹理特征来确定，但并非所有这种特征都是CSDH。气颅的像素区域非常小，如图2所示，每个扫描中的像素数量为15到86个。专家总是需要放大CT图像来观察这个病变，这非常费力和耗时。所有这些都为网络设计带来了挑战。

III Network Architecture

图4显示了作者基于著名的U-Net提出的网络，该网络包含编码器和解码器路径以及编码器和解码器路径之间的直接跳过连接。编码器方面采用最先进的分类网络ResNeSt50作为 Backbone ，因为它结合了通道注意力和ResNet来丰富特征提取。

通过将通道分组并在它们上执行通道注意来突出重要的通道组并抑制不必要的通道。在本论文中，为了进一步选择性地增强所需特征，不同分辨率的提取特征图首先通过空间注意力进行选择性增强以突出形状，然后通过解码器路径上的插值和通道注意进行聚合。

此外，与原始U-Net中的直接跳过连接不同，本文中的跳过连接由1×1卷积和ReLU组成，以降低较低的复杂性。在这些跳过连接中，为了保留精确的位置信息，最低级特征来自输入图像在两个3×3卷积和ReLU层后的直接跳过。对于解码器路径，本文在修复破坏的通道映射以获得更好的类别和形状预测方面，增加了通道注意力块。

Decoder Side – Channel Attention

通道图可以被视为特定类别的响应。然而，在一系列路径聚合和转换之后，通道图会被破坏。为了解决这个问题，作者在解码器侧添加了通道注意力，使上下文表示更具类特定性。

作者将前一层输出和空间注意力块的输出进行拼接，将结果上采样两倍，然后应用具有全局平均池化的通道注意力模块和两个全连接层（FC）及符号函数。在这个模块中，作者使用Group Normalization (GN)而不是Batch Normalization (BN)以更好地适应GPU内存大小并获得更好的性能。

Path Aggregation With Space Attention

对于 Head CT扫描中的颅内出血，位置信息是一个重要特征，因为某些类型的病变只发生在某些区域。某些病变的形状和大小差异很大，这使得模型很难正确预测。此外，一些病变非常小，其空间细节在级联卷积和变换过程中很容易丢失。

为了解决这些问题，本文采用空间注意力[18]如图5所示，以自动学习关注特定病变区域，从而消除后处理的必要性。这个模块通过将通道注意力和相邻不同分辨率特征合并，并对其应用空间注意力来聚合不同层级的特征图。在这里，作者采用Group Normalization (GN)和实例归一化(IN)代替Batch Normalization (BN)，以获得更好的输入特定结果。

IV Result

Experimental Setting

该模型使用连续的三张CT图像作为输入，通过结合相邻信息更准确地预测中心图像。输入进一步通过许多自增强技术（如随机旋转、裁剪、翻转、改变对比度、亮度和饱和度）进行增强，以避免由于数据过少而导致的过拟合。模型在四个NVIDIA RTX-2080Ti GPU上进行训练，使用AdamW优化器，循环学习率，混合精度。

此外，损失函数采用类加权Focal Loss来缓解数据不平衡问题。使用Dice得分系数作为评估指标，并将数据集分为训练（78%）、验证（7.5%）和测试集（14.5%）。

值得注意的是，为了避免对准确性的过高估计，作者根据不同的患者案例而不是随机选择切片图像来划分数据集，因为同一患者案例的分布相似。

Result

表1显示了评估结果。所提出的网络（RSU + SC）在大多数病变病例中的准确率最高，在其他病变病例中也接近最佳，这表明作者对这些具有挑战性的病变类型的网络具有强大的泛化能力。

与表1中的先前网络相比，由于SDH、EDH和SAH等特定病变具有高度相似的特征，它们学习起来很困难。预测结果如图1和6所示。

对于注意力机制，空间注意力和通道注意力各自具有优势。在路径聚合部分使用空间注意力的原因是它们通常发生在某些特定的区域。空间注意力有助于提取位置信息并使形状更准确。在解码器部分使用通道注意力可以改善大多数情况的检测，因为它减少了错误的识别机会并使预测形状更准确。

最后，所提出的网络（RSU + SC）结合了空间注意力和通道注意力的优点。它解决了由路径聚合引起通道图破坏的问题。与基准ResNet50 + U-Net相比，作者提出的模型中每个病变的Dice系数改进分别为2.25%，28.6%，25.1%，31.8%，6.1%，1%对于ICH、SDH、SAH、EDH、CSDH、气颅、脑室内出血。

V Conclusion

本文提出了一种全注意力U-Net来解决CT扫描上颅内出血的分割问题。所提出的网络使用通道注意来突出特定类别的特征，使用空间注意来突出与疾病相关的特征。注意力机制被集成到U-Net的两侧，以在不增加太多参数的情况下获得最大的好处。评估结果显示与基准相比提高了31.8%，并优于其他分割网络。

参考

[1].All Attention U-NET for Semantic Segmentation of Intracranial Hemorrhages In Head CT Images

点击上方卡片，关注「AI视界引擎」公众号