LV-UNet 轻量级的图像分割模型！

未来先知

发布于 2024-10-21 11:11:34

3380

尽管在大规模AI模型在计算机视觉领域取得了进步,优化挑战,Transformer模型的复杂性,计算限制以及实际应用的要求需要更简单的模型架构,尤其是在移动医疗设备上,这些设备需要轻量级且可部署的模型以实现实时性能。然而,一些当前的轻量级模型在不同的数据集上表现出 poor 鲁棒性,这阻碍了它们的更广泛应用。本文提出一个轻量级和基础模型,称为 LV-UNet,有效利用预训练 MobileNetv3-Large 模型并引入可熔合模块。它可以使用改进的深度训练策略进行训练,并在推理过程中切换到部署模式,减少参数数量和计算负载。实验在 ISIC 2016,BUSI,CVC-ClinicDB,CVC-ColonDB 和 Kvar-SEG 数据集上进行,与最先进的模型和经典模型相比,取得了更好的性能。

I Introduction

医学图像分割，旨在准确界定解剖结构或异常，在疾病诊断和治疗计划中至关重要。与传统方法（如依赖于手工或半自动方法）相比，计算机辅助技术可以更高效。随着深度学习的发展，这些方法利用大型标注数据集和设计良好的神经网络学习复杂的图像表示和自动推理像素级标签，使得医学图像分割得到了革新。

最近，医学成像解决方案正在从实验室环境转向床边环境。称为“床边成像”的方法涉及直接在患者身边进行测试和分析，旨在改善患者诊疗，这要求医疗 AI 模型更具实时性和轻便性。例如，床边超声（POCUS）的应用包括妇产科、妇科、急诊医学、传染病、心脏病、腹部、肺部和血管疾病。在资源有限的环境下，初级医疗水平上具有扩展诊断成像能力的 POCUS AI 系统有望带来有意义的健康后果 [1]。

为 POCUS 定制 AI 系统需要考虑轻量级设计、易部署和实时性能。例如，Dulmage 等人开发了一种名为 "Point-of-care and Real-time AI System" 的床边和实时 AI 系统，支持临床医生诊断各种皮肤疾病 [2]。

在资源受限的设备上的轻量级和快速分割至关重要，因为它确保 AI 系统可以快速、准确地划分出感兴趣区域，从而促进资源受限设备上的高效操作并提供精确、及时的结果。

UNet [3] 是一种经典的医学图像分割模型，包括编码器和解码器以及跳过连接，相比之前的尝试有了显著的改进。UNet 的变体，如 Attention-UNet [4]，U-Net++ [5]，ResUnet [6]，Unet 3+ [7] 等也取得了巨大的成功。

由于近年来变形器在视觉任务中的出色表现，Transformer 基于的 U-Net 变体，如 Trans-UNet [14]，Swin-UNet [15]，Unetr [16] 等已得到了更广泛的认可，并显示出显著的有效性。然而，这些模型涉及大量的计算开销和参数，在床边应用中难以使用。

因此，越来越多的努力用于设计轻量级医学图像分割模型。尽管现有的轻量级医学图像分割模型已经取得了较大的进展，但在不同数据集上的鲁棒性仍是一个挑战。如何通过设计具有较低计算开销、较少参数、更快的推理时间和保持各种数据集性能的模型来解决这些问题仍是一个难题。

本文的模型设计基于以下几点：

1）在分割模型的编码器中利用预训练的轻量级模型是至关重要的，确保模型对不同数据集具有鲁棒性；

2）将预训练和扩展模块结合形成设计轻量级医学图像分割模型的简单方法；

3）引入融合模块可以在推理阶段进一步减少参数数量和计算负载。

作者提出了一种名为 LV-UNet 的 UNet 变体，具有轻量级和基本模型设计，包含来自 MobileNetv3-Large 的预训练权重和可配置的扩展模块，并采用 VanillaNet 和融合方法的深度训练策略。

通过在 ISIC 2016 [17]，BUSI [18]，CVC-ClinicDB [19]，CVC-ColonDB [16] 和 Kvair-SEG 数据集 [1] 上进行实验，结果表明，与最先进的经典方法相比，所提出的方法具有具有竞争力的性能。

II Related Works

Lightweight medical image segmentation models

针对医学图像分割的轻量级模型设计已经进行了大量的研究，主要关注U-Net的各种变体。Vallanarasu等人[10]提出了一种基于卷积自回归模型（MLP）的网络UNeXt，包含一个初始卷积阶段和一个潜在MLP阶段，其中MLP块使用了分词化的MLP块，可以投影卷积特征，从而减少参数数量和计算复杂性，提高分割性能，通过将输入通道转换为学习局部依赖性。

Ruan等人[11]提出了一种名为MAL-UNet的模型，该模型集成了一些不同的模块：DGA用于全局和局部特征提取，IEA用于数据集描述和样本连接增强，CAB用于生成多阶段的特征融合，其中包含通道轴自注意力图，以及SAB用于在多阶段特征生成空间轴自注意力图。

Ruan等人[12]提出的EGE-UNet，集成了Group multi-axis Hadamard Product Attention（GHPA）和Group Aggregation Bridge（GAB）模块，以有效提取多样病理信息和融合多尺度特征，显著减少模型大小。

刘等人[13]提出的旋转UNet，将CNN与MLP相结合，以通过有效捕捉局部特征和长程依赖性提高医学图像分割。李等人[12]引入了一种名为U-KAN的增强版U-Net，该模型集成Kolmogorov-Arnold网络（KANs）[14]，以提高医学图像分割和扩散模型的准确性和可解释性，与传统U-Net方法相比，降低了计算成本并取得了优越性能。

VanillaNet

"Vanilla"通常用来形容某物简单、普通或无需装饰，没有多余的部分。追求设计简约的VanillaNet[11]采用了一种简单的架构，优先考虑计算效率和信息保留。通过结合简约操作、批量归一化和激活函数，VanillaNet在不过分复杂的情况下实现了 competitive 的性能。深度训练策略和系列提示激活函数是保持浅层网络非线性问题的关键。

MobileNetw3

MobileNetw3是用于移动和边缘设备的最新一代卷积神经网络，它既拥有高性能又具有高效性。它基于MobileNetw1和MobileNetw2 [12]，通过集成先进的技巧，如逐点（Depthwise）可分隔卷积、倒置残差带有线性瓶颈（Inverted Residuals with Linear Bottlenecks）和Squeeze-and-Excitation（SE）模块来实现。此外，它引入了H-Swish激活函数，以提高计算效率和性能。MobileNetw3使用神经架构搜索（NAS）来优化其结构，以适应不同的性能和效率需求，从而产生两种变体：MobileNetw3-Large以获得更高的准确性和MobileNetw3-Small，用于资源受限的环境。

此外，它还采用了训练技术，如Dropout、标签平滑（Label Smoothing）和Mixup，以获得更好的泛化能力。

III Motivation

The Role of Pre-training

在大型数据集上进行预训练，可以使模型学习到丰富的、概括化的特征表示，从而显著提高它们在各种任务上的表现。这些学习的特征可以捕捉到数据中重要的模式和细微之处，这些可能通过在小数据集上或从零开始训练无法轻易获得。轻量级模型具有较少的参数，这可能会限制它们从数据中学习和泛化的能力。

因此，尽管对模型架构进行复杂的改变可能看起来有利，但这些修改通常并不能为抵消使用预训练模型带来的巨大性能提升。另一方面，虽然复杂的模型结构可能在特定的数据集上提供改进，但这种好处通常是相互矛盾的。

这些复杂的设计可能会过度拟合特定的数据特征，导致在不同的或未见过的数据集上的泛化效果不佳。依赖预训练权重是一种更可靠的方法，以确保模型的整体性能和鲁棒性。

Efficient Model Deployment through Reparameterization

在训练阶段，可以使用一个更复杂的模型来确保最优的学习和特征提取。这个模型可能包含冗余的参数，并使用停用（dropout）和批量归一化（batch normalization）等方法来提高泛化性能。训练完成后，模型将进行再参数化以进行推理。这涉及通过压缩和融合参数（如将多个卷积滤波器合并为更少的滤波器）来简化模型。这样的再参数化模型在保持性能的同时，显著减少了计算和内存需求，确保在资源受限的环境中高效部署。

IV Methods

Overall Architecture

模型架构包括编码器、解码器和最后一个输出层，如图1所示。编码器包括来自MobileNetv3-Large的预训练模块，以及可熔合的编码器模块。解码器由可熔合的解码器模块组成。编码器和解码器之间的对应层之间有跳跃连接，为了避免显著增加参数数量，采用元素加法而非连接。每个阶段的几何尺寸分别为。

Pre-trained Modules in Encoder

LV-UNet采用MobileNetv3-Large模型从初始卷积阶段到第九个反向残差模块作为编码模块。

初始卷积阶段接收输入图像并执行基本特征提取，其余模块包含反向残差模块和逐点分离卷积层，具有各自不同的卷积核和激活函数，逐步提取更多抽象特征。作者将MobileNetv3-Large模型中的15个反向残差模块命名为到。这些反向残差模块分为如下三组：和作为第一组 - I组，，和作为第二组 - II组，，，和作为第三组 - III组。

作者没有使用完整的MobileNetv3-Large（不包含全连接层）以允许后续扩展模块自由设计。如果作者在之后包括反向残差模块到，图像下采样操作的数量（因此图像大小约束）将限制融合模块。

Fusible Modules

Iv-C1 Architecture

可熔丝编码块的训练模式包含一个使用核，步长为1，批标准化操作，尖点ReLU函数，卷积层和使用核的卷积层，一个池化层和一个设计好的激活层。编码块的输出特征的分辨率是输入的一半。在部署模式下，该块仅包含卷积层，池化层和使用设计好的激活层。

可熔丝解码块类似，其中池化操作被替换为上采样操作。解码块的输出特征的分辨率是输入的两倍。如图2所示，训练和部署模式下的可熔丝块的架构。在"深度训练策略与部署模式"部分将介绍在推理阶段将训练模式转换为部署模式的方法。

Iv-C2 Non-Linear Activation Layer

LV-UNet中的非线性激活函数遵循VanillaNet的设计。考虑一个输入特征张量，记作，其中，和分别表示高度、宽度和通道数。激活函数的公式如下：

其中表示输入张量在位置处的激活函数输出，的范围为到，的范围为到，的范围为到。该函数将邻域位置的输入聚合在一个定义在到之间的范围内。

权重控制每个输入的贡献，表示应用于每个输入的激活函数，通常是非线性函数如ReLU。此外，表示每个通道的偏置项。非线性激活函数之后是卷积和批量归一化操作。

利用这种表示形式，LV-UNet可以捕获并合并相邻位置的上下文信息，使模型能够学习和表示输入数据中的复杂模式和非线性关系。

Iv-C3 Deep Training Strategy and Deployment Mode

原论文中的VanillaNet采用了一种深度训练策略，并在LV-UNet的训练中应用到了。在训练过程开始时，采用两个卷积层和激活函数。在训练轮次中，激活函数逐步调整，以模仿一种对齐映射。训练结束时，这两个卷积可以无缝地结合成单个卷积，从而减少了推理时间。作者使用Leaky ReLU函数来控制它：

其中是当前轮次，是深度训练轮次数。在训练开始时，激活函数是ReLU函数，而在训练结束时，它是一种对齐映射函数。

非线性逐渐降低。与原文[22]中定义为的设计不同，作者设计的方法可以使从线性到非线性的过渡呈现出开始慢，然后加速的趋势，这有助于训练收敛。

遵循VanillaNet训练方法，将单个卷积层weight和bias矩阵（来源于批标准化层和其前面的卷积层，在可合并块中）合并，得到：

其中，是卷积核的权重和偏置矩阵，，是批标准化中的尺度，偏移，均值和方差，下标表示-th输出通道的值。

然后，可以将fusible块中的两个1x1卷积层合并为一个卷积层：

其中表示卷积操作，代表矩阵乘法，和是两个卷积层的权重矩阵，是通过im2col操作从输入转换成一个对应核形状的矩阵。

通过使用深度训练策略，LV-UNet可以在推理阶段切换到部署模式，这可以大大降低参数和计算要求。

V Experiments

Datasets

V-A1 Isic2016

ISIC 2016 [11] 是一个在2016年国际生物医学图像研讨会（ISBI）上进行的皮肤病变分析朝着黑色素瘤检测挑战的数据集。在本论文中，作者使用了病变分段数据集，包含原始图像和二进制病变 Mask 。训练集中有900对原始图像和相应的 Mask 。测试集中有379对图像和 Mask 。训练集为Iv-A2 Business。

乳腺超声图像数据集（Dataset BUSI）[12] 是收集自25至75岁女性之间的乳腺超声图像的集合。该数据集包括正常、良性和恶性乳腺超声图像，以及相应的 Mask 。

与[16]中的设置不同，作者包括所有三个类别在数据集中，并且测试集按照0.2的比例与随机状态42进行划分。然后剩下的数据集以0.2的比例与随机状态42划分训练集和验证集。

Iv-A3 CVC-ClinicDB

CVC-ClinicDB [13] 是一个从结肠镜视频提取的多息肉帧的图像数据库。该数据集包含 612 对图像及其对应的 Mask 。测试集按比例 0.2 分割，随机状态为 42。然后剩余的数据集按比例 0.2 分割为训练集和验证集，随机状态为 42。

Iv-A4 CVC-ColonDB

CVC-ColonDB 数据集 [14] 包括380张图像，每个图像的分辨率为574×500像素。这些图像是从15段结肠镜视频（共12,000张图像）中选取的，其中只有这380张图像进行了标注。测试集以0.2的比例与随机状态42分片。然后，其余的数据集以0.2的比例与随机状态42分成了训练集和验证集。

Iv-A5 Kvasir-SEG

Kvasir-SEG数据库[15]包括1000个多面体图像及其对应的真实 Mask ，分辨率从332487像素到19201072像素不等。测试集以0.2的比例进行划分，随机状态为42。然后，其余的数据集以0.2的比例划分为训练集和验证集，随机状态为42。

Implementation Details

所有实验均在具有32 GB内存的PG500-216(V-100)上执行。总的训练轮次为300，训练、验证和测试中的批量大小均为8。输入图像的分辨率被重新缩放到256×256。使用的优化器是ADAM[24]。初始学习率是0.001，并使用CosineAnnealingLR[25]调度器。最小学习率是0.00001。数据增强方法包括仅顺时针旋转90度，随机翻转和数据归一化。验证集上的评估指标是病变的IOU值。使用的损失函数是一种结合二进制交叉熵（BCE）损失和 dice 损失的混合损失：

作者还训练、验证和测试了UNet、Attention UNet、UNet++、UNext、MLUNet、EGE-UNet、Rolling-UNet-M、U-KAN和MobileNetv3-UNet（作为 Backbone 的网络结构UNet）。

UNet在每个阶段的UNet中使用双卷积层。UNet、Attention UNet和UNet++的每个阶段尺寸分别为[32, 64, 128, 256, 512]，而UNext、MLUNet、EGE-UNet、Rolling-UNet-M和U-KAN使用其论文中的原始设计。

MobileNetv3-UNet在解码器中的每个块上将每个编码器阶段的输出与添加到MobileNetv3 Backbone 结构中的块连接。

在其他实验中，作者选择在验证集上表现最佳的模型在测试集上进行推理。然而，Leaky ReLU只在训练结束时转型为恒等映射函数。

因此，作者选择最后一个周期的模型在测试集上进行推理（对于使用深度训练战略的LV-UNet）。

Results

表1中的对比实验结果可以看出，LV-UNet相较于 Baseline 具有竞争性能。LV-UNet的部署模式在降低参数的同时仍能实现良好的性能。表2中，作者将不同模型的参数数和GFLOPs进行了比较，结果显示LV-UNet及其部署模式的大小和计算复杂度非常适合边缘设备和点到关怀场景。图3中给出了部分分割结果的示例。

在表1中，作者可以看到LV-UNet相较于 Baseline 具有出色的性能。尽管LV-UNet的部署模式在降低参数方面做出了很大贡献，但它仍然在很大程度上维护了良好的性能。在表2中，作者将不同模型的参数数量和GFLOPs进行了比较，这表明LV-UNet及其部署模式的大小和计算复杂度完美地适合边缘设备和点到关怀场景使用。图3展示了部分分割结果的样例。