UNet v2开源啦 | 参数更少，猛涨20个mIoU，甚至还可以更多

点击下方卡片，关注「集智书童」公众号

在本文中，作者介绍了一种新的、鲁棒且高效的U-Net变体，用于医学图像分割。它的目标是同时增强低级特征中的语义信息注入，并使用更细的细节精炼高级特征。对于输入图像，作者首先使用深度神经网络编码器提取多级特征。接下来，通过将高级特征的语义信息与低级特征的细粒度信息通过哈达马乘积相结合，增强每级特征图。作者的新跳过连接使所有级别的特征都具有丰富的语义特性和复杂的细节。改进的特征随后被传输到解码器进行进一步处理和分割。作者的方法可以无缝集成到任何编码器-解码器网络中。

作者在几个公共的医学图像分割数据集上评估本文的方法，包括皮肤病变分割和息肉分割，实验结果表明，与最先进的方法相比，新方法具有更高的分割精度，同时保持了内存和计算效率。

代码：https://github.com/yaoppeng/U-Net_v2

1 Introduction

随着现代深度神经网络的进步，语义图像分割取得了显著的进展。语义图像分割的典型范式涉及一个编码器-解码器网络，其中包含跳过连接。在这个框架中，编码器从输入图像中提取分层的和抽象的特征，而解码器则使用编码器生成的特征图，以像素为单位重建分割掩膜或图，并对输入图像中的每个像素分配一个类标签。已经进行了一系列研究，以将全局信息集成到特征图，并增强多尺度特征，从而在分割性能方面取得了显著的改进。

在医学图像分析领域，准确的图像分割在计算机辅助诊断和分析中起着关键作用。U-Net最初用于医学图像分割，利用跳过连接在每个级别的编码器和解码器之间进行连接。跳过连接使解码器能够访问来自较早的编码器阶段的特征，因此同时保留了高级语义信息和细粒度空间细节。这种方法便于在医学图像中精确划分对象边界和提取小结构。此外，采用密集连接机制通过将所有级别的所有阶段的特征连接在一起来减少编码器和解码器之间特征的差异。设计了一种机制，通过将来自较高和较低级别的不同尺度的特征连接在一起来增强特征。

然而，基于U-Net的模型中的这些连接可能并不足以有效地集成低级和高级特征。例如，在ResNet中，一个深度神经网络被构建为一个多个浅层网络的集成，并且明确添加的残差连接表明，即使在一个亿像素图像数据集上进行训练，网络也很难学习身份映射函数。

关于编码器提取的特征，低级特征通常保留更多的细节，但缺乏足够的语义信息和可能包含不希望的噪声。相比之下，高级特征包含更多的语义信息，但由于分辨率显著降低，因此缺乏精确的细节（例如，对象边界）。简单地将特征通过串联进行融合将严重依赖于网络的学习能力，这通常与训练数据集的大小成正比。这是一个具有挑战性的问题，尤其是在医学影像的背景下，通常受到数据的限制。通过在多个级别之间通过密集连接串联低级和高级特征，可能限制不同级别信息的贡献，并可能引入噪声。另一方面，尽管新增的卷积操作并未显着增加参数数量，但GPU内存消耗将上升，因为所有中间特征图和相应的梯度都必须在正向传播和反向梯度计算中进行存储。这导致GPU内存使用和浮点运算（FLOPs）的增加。

在[8]中，反向注意力被用来明确建立多尺度特征之间的联系。在[9]中，ReLU激活被应用到高级特征，并将激活的特征与低级特征相乘。此外，在[10]中，作者提出分别从CNN和Transformer模型中提取特征，将CNN和Transformer分支在多个级别上的特征组合，以增强特征图。然而，这些方法都较为复杂，且其性能尚不令人满意，因此需要进一步改进。

在本文中，作者提出了一种新的基于U-Net的分割框架，具有简单且高效的跳过连接。作者的模型首先使用CNN或Transformer编码器提取多级特征图。接下来，对于第i级的特征图，作者通过简单的哈达马乘法操作明确注入包含更多语义信息的高级特征和包含更精细细节的低级特征，从而增强第i级特征的语义和细节。随后，优化的特征被传输到解码器进行分辨率重建和分割。作者的方法可以无缝集成到任何编码器-解码器网络中。

作者在两个医学图像分割任务上评估了作者的新方法，即皮肤病变分割和息肉分割，使用了公开可用的数据集。实验结果表明，作者的U-Net v2在这些分割任务上始终优于最先进的方法，同时保持FLOPs和GPU内存效率。

2 Method

Overall Architecture

作者的U-Net v2的整体架构如图1(a)所示，它包括三个主要模块：编码器、SDI（语义和细节注入）模块和解码器。对于一个输入图像，其中，编码器产生个级别的特征。作者用表示第个级别的特征，其中。这些收集的特征然后被传输到SDI模块进行进一步的优化。

Semantics and Detail Infusion (SDI) Module

在使用编码器生成的层次特征图的基础上，作者首先对每个级别的特征应用空间和通道注意力机制。这个过程使得特征可以集成局部空间信息和全局通道信息，具体如下所示：

在这里，表示第个级别的处理后的特征图，和分别表示第个级别中的空间和通道注意力的参数。此外，作者应用一个卷积来将的通道数减少到，其中是一个超参数。这样得到的结果特征图被定义为，其中，其中，和分别表示的宽度、高度和通道数。

接下来，作者需要将优化的特征图发送到解码器。在每个解码器级别中，作者使用作为目标参考。然后，作者调整每个-th级别的特征图大小，以匹配的相同分辨率，具体如下所示：

在这里，，和分别表示自适应平均池化，恒等映射和双线性插值到×的分辨率，其中。然后，一个3×3卷积被应用以平滑每个重采样特征图，具体如下所示：

在这里，θ表示平滑卷积的参数，是在第i个级别中的第j个平滑特征图。

在将所有第i个级别的特征图重新缩放到相同分辨率之后，作者对所有重采样特征图应用逐元素哈达马乘法，以增强第i个级别特征图，同时具有更多的语义信息和更细的细节，具体如下所示：

其中，H(·)表示哈达马乘法（参见图1(b)）。然后，被传递到第i个级别的解码器进行进一步的分辨率重建和分割。

3 Experiments

Datasets

作者使用以下数据集来评估作者的新U-Net v2。

ISIC数据集：用于皮肤病变分割的两个数据集是：ISIC 2017，包含2050张皮肤病变镜像，ISIC 2018，包含2694张皮肤病变镜像。为了公平比较，作者遵循[13]中规定的训练/测试划分策略。

息肉分割数据集：作者使用了五个数据集：Kvasir-SEG，ClinicDB，ColonDB，Endoscene和ETIS。为了公平比较，作者遵循[8]中的训练/测试划分策略。具体而言，ClinicDB中的900张图像和Kvasir-SEG中的548张图像被用作训练集，其余图像被用作测试集。

Experimental Setup

作者在NVIDIA P100 GPU上使用PyTorch进行实验。作者的网络使用Adam优化器进行优化，初始学习率等于0.001，等于0.9，等于0.999。作者使用幂学习率衰减，衰减系数为0.9。最大训练周期设置为300。超参数设置为32。

如同[13]中的方法，作者在ISIC数据集上报告DSC（Dice相似系数）和IoU（交并比）得分。对于息肉数据集，作者报告DSC，IoU和MAE（平均绝对误差）得分。每个实验运行5次，报告平均结果。作者使用Pyramid Vision Transformer（PVT）作为特征提取的编码器。

Results and Analysis

在ISIC数据集上的比较结果如表1所示。如所示，作者提出的U-Net v2在ISIC 2017和ISIC 2018数据集上分别将DSC分数提高了1.44%和2.48%，并将IoU分数提高了2.36%和3.90%。这些改进证明了作者的提出的方法将语义信息和更精细的细节注入每个特征图的有效性。

息肉分割数据集上的比较结果如表2所示。如所示，作者提出的U-Net v2在Kavasir-SEG，ClinicDB，ColonDB和ETIS数据集上分别超过了Poly-PVT，在这些数据集上的DSC分数分别提高了1.1%，0.7%，0.4%和0.3%。这进一步证明了作者在每个级别将语义信息和更精细的细节注入特征图的提出方法的一致有效性。

Ablation Study

作者使用ISIC 2017和ColonDB数据集进行消融研究，以检查作者的U-Net v2的有效性，具体如表3所示。具体而言，作者使用PVT模型作为UNet++的编码器。请注意，当作者的SDI模块被移除时，U-Net v2退化为一个普通的U-Net，具有PVT Backbone。

SC表示SDI模块内的空间和通道注意。从表3可以看出，与没有SDI的U-Net v2（即具有PVT编码器的U-Net）相比，UNet++表现出轻微的性能降低。这种下降可能归因于由密集连接生成的多级特征的简单串联，这可能会使模型困惑并引入噪声。表3证明了SDI模块对整体性能的贡献最大，这突出了作者提出的跳过连接（即SDI）始终产生性能改进的一致性。

Qualitative Results

一些在ISIC 2017数据集上的定性示例如图2所示，这些示例证明了作者的U-Net v2能够将语义信息和更精细的细节集成到每个特征图，因此作者的分割模型可以捕捉到物体边界的更细微细节。

Computation, GPU Memory, and Inference Time

为了检查作者的U-Net v2的计算复杂性、GPU内存使用和推理时间，作者在表4中报告了作者的方法U-Net、UNet++的参数、GPU内存使用、FLOPs和FPS。所有实验都使用float32作为数据类型，导致每个变量使用4B的内存。GPU内存使用记录了在正向/反向传播过程中存储的参数和中间变量的尺寸。表示输入图像的大小。所有测试都在NVIDIA P100 GPU上进行。

在表4中，作者可以观察到UNet++引入了更多的参数，并且由于在密集正向过程中存储了中间变量（例如，特征图），其GPU内存使用更大。通常，这些中间变量比参数消耗更多的GPU内存。此外，U-Net v2的FLOPs和FPS也优于UNet++。与U-Net（PVT）相比，作者的U-Net v2的FPS减少是有限的。

4 Conclusions

一种新的U-Net变体U-Net v2被引入，该变体具有改进医学图像分割的新颖和简单跳过连接设计。这种设计通过使用哈达马乘法将语义信息从高级特征和从低级特征中获取的更精细的细节集成到由编码器生成的特征图的每个级别中。在Skin Lesion和Polyp Segmentation数据集上的实验验证了作者的U-Net v2的有效性。复杂性分析表明，U-Net v2在FLOPs和GPU内存使用方面也具有效率。

参考

[1]. U-NET V2: RETHINKING THE SKIP CONNECTIONS OF U-NET FOR MEDICAL IMAGE SEGMENTATION.

点击上方卡片，关注「集智书童」公众号