U-Net 架构的演进,结合领域分解与通信网络的超高分辨率图像分割新策略 !

2024-08-08 11:27:08 浏览数 (2)

在本次研究中,针对超高分辨率图像分割中出现的空间信息丢失和计算效率低下等问题,提出了一种将编码器-解码器结构与领域分解策略相结合的创新方法。 具体来说,本研究引入了一种基于领域分解的U-Net架构(DDU-Net),它能将输入图像分割成可以独立在不同设备上处理的不重叠块。此外,还增加了一个通信网络,以便于块与块之间的信息交换,从而增强对空间上下文的理解。在合成数据集上进行的实验验证旨在衡量通信网络的有效性。 随后,在DeepGlobe土地覆盖分类数据集上进行了性能测试,该数据集作为现实世界的基准数据集。实验结果表明,对于划分为16×16不重叠子图像的图像,包含块间通信的方法比不包含通信的相同网络能获得高出2-3%的交并比(IoU)得分。包含通信的网络性能与在完整图像上训练的基准U-Net相当,这表明本模型在保持空间上下文的同时,为超高分辨率图像的分割提供了一种有效的解决方案。 相关代码可在https://github.com/corne00/HiRes-Seg-CNN获取。

1 Introduction

在本文中,大部分计算机视觉领域的深度学习模型专注于低分辨率的2D和3D图像,通常是256×256像素或更小。然而,随着高分辨率图像数据集的广泛应用,由于单个GPU的内存限制,带来了新的挑战,尤其是在图像语义分割等内存密集型任务中。语义分割是将输入图像中的像素分类到不同的、不重叠的语义类别中的计算机视觉任务。超高分辨率图像分割在卫星图像中的目标分割、金属表面缺陷检测以及计算机辅助医学诊断等多样化领域具有重要意义。尽管深度卷积神经网络(CNNs)在图像分割方面取得了显著的成功,但大多数模型由于对内存的高要求,不适用于高分辨率图像的模型训练和推断。

以亚毫米分辨率的计算机断层扫描(CT)为例,其体素图像数据通常具有512×512×512体素的分辨率。即使使用半精度浮点数和适中的批量大小8,仅用一层卷积神经网络(64个滤波器)处理此类图像,也需要超过137GB的GPU/TPU内存,如文献[27]所示。使用传统策略,如降采样或块裁剪来处理如此高分辨率的输入,通常会导致详细信息或空间上下文的丢失,从而降低分割精度。

本文提出了一种基于著名的U-Net网络的新模型架构,该架构最初是为生物医学图像分割任务而开发的。U-Net是一种具有瓶颈结构和编码器与解码器路径间跳跃连接的卷积神经网络(CNN)。这一基础架构及其变体在各种语义分割和其他图像到图像的任务中展示了卓越的、最先进的准确性。然而,其较高的内存需求限制了在内存有限的计算设备上应用于高分辨率图像的能力。训练过程中的高内存需求来自于在正向传播过程中获得的中间特征图的存储,以便在随后的反向传播过程中使用。特别是在处理高分辨率输入时,这导致网络的第一层和最后几层出现高维特征图。

作者提出的模型结合了U-Net架构和受领域分解方法(DDMs)启发的分而治之策略,以应对计算设备内存限制。之前的并行化策略通过将图像分解为子图像(子域)来进行处理,但在每次卷积之前通过通信边缘或冗余计算提供子图像间的全局上下文信息。我们的方法同样分解图像,但通信明确限制在U-Net架构的瓶颈部分。这最小化了通信开销,同时保留了必要的上下文信息。

作者的主要贡献总结如下:

  • 提出了一种结合U-Net架构与领域分解策略的新方法,以高效地分割超高分辨率图像,同时保持空间上下文。
  • 展示了通信网络,即我们方法的重要组成部分,可用于不同子图像间的信息交换,增强对空间上下文的理解,而不会带来显著的计算开销和额外的通信与内存成本。
  • 通过在合成和真实图像数据集上评估我们的架构,证明了与基准U-Net模型相比具有竞争力的分割性能。即使训练仅限于可处理的最大图像部分,我们的方法仍然具有可扩展性。

2 Related Work

在本节中,作者将讨论U-Net网络,这是作者所提出方法的基础,并介绍感受野这一关键概念,它对于理解卷积神经网络(CNNs)至关重要。此外,作者回顾了以往解决U-Net及基于CNN的语义分割模型在内存挑战上的尝试;作者讨论了这些方法的优点和不足。

最后,作者引入域分解这一概念,作为一种自然处理内存限制的方法,并强调了一些先前的研究,这些研究将此方法应用于分类和分割任务。

The U-Net

本文提出的新的分割方法基于U-Net架构[49]。自该架构发表以来,已经发展出了许多变体和扩展,例如3D U-Net[10]、UNet [59]、Attention U-Net[45]和ResUNet-a[15]。作者假设读者熟悉CNN[35]的基本概念以及这些模型的构建模块——卷积层和池化层。关于这些概念的详细解释,作者参考了该主题丰富的文献资料,例如[21,第9章]或[6,第10章]。

U-Net架构[49],如图1所示,由两条路径组成:收缩路径(编码器路径)和扩张路径(解码器路径)。收缩路径包括重复的块,每个块由两个连续的卷积组成,之后是ReLU激活函数和最大池化层。相反,扩张路径使用转置卷积对特征图进行上采样。随后,将收缩路径中相应层的特征图进行裁剪(如果上采样路径舍弃了边界像素),并与上采样的特征图进行拼接。拼接之后是两个连续的卷积和一个ReLU激活函数。

在最后阶段,通过卷积将特征图的数量减少到所需的类别数,生成分割图像。收缩路径(捕捉全局上下文和特征)与扩张路径(实现精确的细粒度定位)之间的跳跃连接,是U-Net与其他相关网络架构(例如[38]中提出的全卷积网络(FCN))的区别之处。需要注意的是,最初提出的U-Net返回一个裁剪后的分割 Mask (由于每次卷积都会丢失边界像素)。然而,U-Net的较新实现通常在边界处使用零填充,以获得与输入尺寸相等的输出尺寸,从而保持了输入尺寸并无需裁剪;参见[7,58,4]。

The Receptive Field

理解CNN的一个重要概念是它们的_感受野_,或称_视野_。CNN的感受野是指输入图像中影响CNN输出中一个像素的区域。与每个输入神经元都与每个输出神经元相连的全连接神经网络不同,CNN的输出依赖于输入层的特定区域。在CNN的每一层中,随着作者深入网络,感受野会扩大,使得更高层次的特征能够受到输入图像更大部分的影响。

感受野的理论大小可以通过实验和理论分析来研究,如文献[48]中详细讨论的。感受野的大小对CNN的预测性能有重要影响。如果感受野太小,网络可能无法捕捉到所有相关信息,导致预测效果不佳。对于像分割这样的任务,感受野应该足够大,以包含输入图像中所有对预测正确像素类别相关的像素。请注意,感受野的最佳大小取决于网络架构以及训练网络的数据。

在实际应用中,CNN的感受野大小受多种因素影响,包括卷积层的数量、卷积核的大小以及池化层的存在;参见[39]。增加网络的深度和核的大小可以扩大感受野,使网络能够从输入图像中捕捉到更多上下文。然而,这也增加了计算复杂性和过拟合的风险。

Overcoming U-Net Memory Limits

大多数U-Net变体共享的一个问题是,由于在训练期间需要存储中间特征图,导致内存需求很大,这使得该模型不适用于有限内存计算设备上的高分辨率应用[2]。克服这些限制的一个自然方法是使用数据并行和/或模型并行,将特征图的内存负载分配到不同的计算设备上。几种并行化的U-Net形式采用了并行化来减少每个设备的内存占用。在[27]和[56]中,作者实现了一种空间划分技术,将卷积层的输入和输出分解为较小的不重叠子图像。在每个卷积操作之前,设备之间交换具有卷积核大小一半的特征图边缘区域。尽管这种方法引入了通信开销,但它有效地跨设备分配内存,而不改变U-Net的基本架构。然而,作者的方法涉及到对U-Net架构的修改,以进一步优化性能并将通信限制在架构的瓶 Neck 分。

文献[50]的作者将图像划分为重叠的子图像,重叠大小由感受野大小决定。由于冗余计算(对于划分为子图像的像素图像以及具有感受野大小的U-Net,计算冗余为),这种方法允许完全并行执行前向和后向传递。[57]采用了类似的策略,但应用于使用ResNet架构[24]的图像分类而不是图像分割。

这些方法成功地将U-Net划分为可以在并行中训练和评估的方式,即使对于超高分辨率图像数据集也是如此。然而,它们要么涉及在每个卷积操作之前进行边缘通信,这导致通过许多点对点消息的通信开销;要么包含许多冗余计算。

Memory Optimization for CNNs Using Domain Decomposition Approaches

为了提高内存效率,由于计算局部化所带来的内在并行性和可扩展性,分布式域分解方法(DDMs)颇具启发意义。DDMs是(离散化)偏微分方程(PDEs)的有效迭代求解器,通过分而治之的策略将计算域划分为重叠或非重叠的子域,展现出可扩展性。PDE问题在这些子域上被划分为子问题。这使得子域内的计算可以并行执行。尽管在潜在的子域边界上通过平衡的邻居通信确保了全局收敛,以及少量的全局通信。例如,高度可扩展的最新DDMs是重叠Schwarz方法的变体[17, 53]或通过约束的平衡域分解(BDDC)[11, 16]以及有限元裂化和互联——双重主方法(FETI-DP)[19, 20]。

近年来,将DDMs与机器学习(ML)算法相结合的兴趣日益增加,以结合两个领域的优势[25, 32]。到目前为止,大多数研究工作集中于将DDMs与ML结合以解决PDEs。然而,本文尝试将卷积神经网络(CNN)与DDM技术结合用于语义图像分割。值得注意的是,据作者所知,尚未有研究明确基于域分解策略来解决图像分割任务,尽管已经提出了一些与域分解策略相似的方法。现有研究主要关注图像分类,这与语义分割密切相关;但是,由于模型输出的不同,网络架构也有所不同。在这里,作者提供了将DDM策略与ML结合用于图像分类和语义图像分割任务的简明概述。

在[42]中,作者提出了一种图像分割的图像划分方法,该方法降低了网络复杂性并提高了并行化能力。他们将输入图像划分为非重叠的子图像,并随后在这些子图像上训练较小的局部CNN,而不是一个大的全局CNN。这增强了模型对特定图像区域的专门化,并在没有对全局CNN的测试误差产生重大影响的情况下降低了整体复杂性。重要的是,作者完全独立处理子图像,它们之间没有耦合。

在[31]中,将DDM策略应用于训练用于图像分类的CNN-DNN(卷积神经网络-深度神经网络)架构。特别是,作者将输入图像划分为子图像,每个子图像用于训练一个单独的CNN以局部预测类别。随后,一个DNN将这些局部预测聚合为一个全局预测。作者将DNN解释为一个粗问题求解器,它结合了来自局部网络的细粒度信息。

在[22]中,作者将用于图像分类的全局CNN沿其宽度分解为有限数量的较小、独立的局部子网络。通过训练这些局部子网络获得的权重作为随后训练全局网络的初始化,采用迁移学习策略。在[44]中,作者提出了一种估计诸如能量或熵等广延参数的物理启发的神经网络拓扑。他们将域分解为具有“焦点”和“上下文”区域的子域,并根据估计的广延参数的“局部性”进行重叠。不同的子域使用相同的子网络权重(作者的方法也采用此策略),这导致参数数量的减少,提高了并行化能力,并加快了推理速度。最后,在[46]中,提出了一个受DDM启发的分割算法,该算法将输入图像划分为多个重叠的子图像,并在这些子图像上训练一个分割网络,允许并行推理。

作者得出结论,这种方法能更好地分割小目标。

3 Methodology

在本节中,作者介绍了基于域分解的U-Net(DDU-Net)架构以及作者的训练方法。此外,作者还描述了用于测试模型的数据库,并概述了模型训练过程。

Network architecture

3.1.1 Decomposition of the images and masks

图2展示了作者提出的网络架构的示意图。模型以一组(高分辨率)图像及其对应的分割 Mask 数据集为起点,处理尺寸为的2D像素图像(左侧),并为图像中的每个像素输出一个类别的概率分布(右侧)。

遵循DDM策略,作者将输入数据分解为个互不重叠的子图像,其中。具体来说,每幅图像被划分为高度为和宽度为的个子图像,和,使得和。相应的分割 Mask 也以类似方式进行划分。子图像和子 Mask 被分配到计算设备(例如,GPU/TPU)上,以平衡工作负载。

图2:所提出网络架构的示意图。输入图像被划分为子图像,在编码器路径中独立处理。编码后,一定数量的编码特征图被发送到包含通信网络的设备,并通过通信网络进行处理。该网络的输出替代了输入特征图。解码过程也是在没有计算设备间通信的情况下并行进行的。虚线箭头表示跳跃连接。编码器-解码器网络和通信网络的详细架构分别展示在图3和图4中。

在划分之后,子图像在不同的计算设备上由编码器独立、并行地处理。关于编码器和解码器架构的更多细节,请参见1.2小节。从编码器最后一层选择的特征图数量被发送到通信网络(见3.1.3小节),允许子图像之间交换相关上下文。通信网络将所有子图像的编码特征图按原始全分辨率图像中子图像的排列方式拼接并处理它们。通信网络的每个输出特征图再次被划分为个子特征图,然后发送回设备。解码器随后为每个子图像生成子预测,将这些子预测拼接起来,形成整个高分辨率图像的全局预测 Mask 。

3.1.2 Encoder-decoder networks

如3.1.1小节所述,每个计算设备包含一个独立的编码器-解码器网络,允许并行处理子图像。然而,这些网络共享其权重,以确保子图像之间的一致性分割,使它们成为全局编码器-解码器网络的局部克隆。这意味着在训练期间,每次反向传播步骤后都需要同步权重和梯度。

每个局部编码器网络使用几个卷积块处理对应的输入子图像,并在每个卷积块后紧跟一个批量归一化层和一个ReLU激活函数。在训练期间,还应用了一个dropout层以防止过拟合。经过两个卷积块后,进行最大池化以减少特征图的空间尺寸。最深的编码器层产生256个空间粗糙的特征图(该数字随着网络架构的变化而变化),其中一些用于编码器-解码器网络之间的通信。每个编码器网络的输出是256个特征图,具有较小的空间尺寸。这256个特征图中的最后个被引导到通信网络中,通信网络处理它们并返回具有相同尺寸的修改后的特征图,以替换原始输入特征图;参见3.1.3小节。通信后,特征图接着通过扩展路径处理,为每个子图像生成最终的分割子 Mask 。这些子 Mask 被连接起来,形成覆盖整个全分辨率图像的完整 Mask 。

3.1.3 Communication network

在本论文中,通信网络由三层卷积层组成。DDU-Net架构并不局限于这种特定的网络配置。根据特定应用调整层的数量、膨胀度和核大小等,可能会增强模型的感受野大小。作者选择通信网络为全卷积网络,以确保能够适应任意输入尺寸,同时可调整子图像的数量。

通信网络接收来自每个子图像最深编码器层的个特征图作为输入;参见3.1.2小节。这些特征图沿着高度和宽度维度进行拼接,与全分辨率输入图像中子图像的位置相对应;参见图2。通信网络处理这个拼接后的输入,并生成与输入相同空间尺寸的特征图。这些特征图沿着高度和宽度再次分割成个子特征图,然后送回到相应的编码器-解码器网络进行进一步处理;它们替换了由编码器生成的原始特征图,并作为通信网络的输入。

Notation

在本论文的其余部分,作者采用一种命名约定来比较基于领域分解的U-Net的不同模型配置,该命名约定包含了关键的参数。每个模型表示为DDU-Net,其中:

  • () 表示编码器-解码器网络中上采样和下采样块的数目。
  • () 指定由通信网络处理特征图的数目。
  • (Y或N) 表示通信特征图之间是否启用通信()或禁用通信()。

需要注意的是,如果禁用通信(),但通信网络处理特征图的数目非零,这意味着个特征图会独立地通过通信网络发送,不进行拼接。这实际上等同于在瓶颈层具有额外卷积操作的 Baseline U-Net架构。

例如,DDU-Net指的是一个模型,其编码器-解码器网络的深度为,个特征图被发送到通信网络,且通信是启用的。而DDU-Net在子图像上的操作等同于 Baseline U-Net架构。在这种情况下,DDU-Net,其中,指的是在瓶颈层具有额外卷积操作的 Baseline U-Net。此外,当时,无论通信是否启用(Y)都没有关系,因为通信网络使用了个特征图。

Model training

损失函数。在网络的训练过程中,作者采用了dice损失函数,因为该函数解决了图像语义分割中的类别不平衡问题;参见[41, 28]。Dice损失定义如下:

在这里,表示类别数量,是一个批次中的像素总数,是像素属于类别的预测概率,通过将softmax函数应用于模型的输出来得到,从而使所有输出逻辑值处于范围内。此外,表示像素属于类别的真实概率,由于已知真实 Mask ,其值限制为,而作为一个小的数值稳定性常数(以避免除以零),在本文中设置为。

参数优化。网络的权重采用He等人[23]提出的方法进行初始化。作者使用Adam优化器[30],动量参数,并采用平台学习率衰减策略,当验证损失不再降低时,学习率以的因子衰减,保持选择的轮次耐心。初始学习率、批量大小和早停准则根据任务和可用内存进行选择,具体将在结果部分提供。

训练过程。在训练DDU-Net架构时,作者首先用共享权重初始化编码器-解码器网络,并将此网络的副本分布到所有设备上。通信的粗略网络在一个选定的设备上初始化,由于通信特征图的尺寸较小,可以选择与其中一个编码器-解码器网络位于同一设备上。在前向传播过程中,每个设备独立计算编码后的特征图。在用通信网络处理选定数量的特征图后,特征图的解码再次完全并行进行。通过编码器和解码器的反向传播可以在没有局部编码器-解码器网络之间依赖的情况下并行进行。然而,在通过解码器反向传播之后,通信网络的反向传播是必要的。在反向传播之后,梯度在主设备上集中累积,并用于更新权重,确保所有设备上的权重更新一致。更新后的权重广播到其他设备进行同步。

实现。该实现使用了PyTorch[47](版本1.12.0),这是一个开源的机器学习库。作者在荷兰代尔夫特理工大学的DelftBlue超级计算机[12]上进行训练和测试,并使用NVIDIA Tesla V100S GPU进行训练,该GPU具有 GB的内存。

Datasets

为了测试作者的模型,作者使用了两个不同的图像数据集:

  1. 一个合成生成的数据集,旨在测试通信网络的能力;
  2. 一个真实图像语义分割数据集,用于多类别土地覆盖分割,以评估所提出模型在分割质量以及内存效率方面的有效性。

3.4.1 Synthetic dataset

在作者的方法中,仅在编码器-解码器网络之间交换具有低空间分辨率的深度特征图。这与先前涉及在每个U-Net层交换一系列特征图的U-Net并行化方法有显著不同。为了评估这些低分辨率特征图能够捕获的空间上下文水平,作者设计了一个合成数据集。该数据集由尺寸为像素的单通道灰度图像组成,其中。对于每个,作者生成了个训练图像,个测试图像和个验证图像,总计合成了个图像。这种设计允许将图像分解为个像素的子图像;图5展示了时的子图像示例。

这个灰度数据集中的每个图像都显示了一个黑色背景,上面有两个随机放置的白色圆圈,每个圆圈的半径为4像素,完全位于一个子图像内,使得子图像边界不会与圆圈相交。对应的 Mask 与图像相似,但增加了一个像素类别:在两个圆圈中心之间绘制一条线段。图5展示了两个图像及其对应 Mask 的示例。

当使用DDU-Net处理此数据集中的图像时,连接两个圆圈的线段的分割完全依赖于通信网络。由于这两个子图像是独立处理的,仅通过通信网络连接,分割的有效性直接反映了通信网络准确传递全局信息的能力。

3.4.2 DeepGlobe land cover classification dataset

DeepGlobe土地覆盖分类数据集[13]是一个针对土地覆盖类型的语义分割数据集。该数据集包含803张高分辨率(像素)的标注卫星图像,分为7个类别:城市、农业、草原、森林、水域、荒地和未知。这些图像具有像素的分辨率,总面积为。除了图像的高分辨率之外,分割这个数据集还面临着大类不平衡的挑战;见表1。在图6中,作者展示了数据集中的两张示例图像及其 Mask 。

其他挑战包括图像数量有限、不精确的 GT 情况以及多尺度相关上下文的存在;参见[37, 29]。为了说明这一点,请注意树木可能存在于各种土地类型中,如城市、草原或森林区域。因此,网络需要从树周围的更广阔区域整合上下文信息,以准确预测树像素的正确类别。

Evaluation Metrics

在训练模型之后,作者分析了结果。由于类间存在较大的不平衡,作者采用了类平均交并比(IoU)得分作为评价指标,这是在引入DeepGlobe数据集[13]的论文中建议的。对于一个包含张图像的数据集,类别的IoU得分定义为:

其中:

  • 是第张图像中正确预测为类别的像素数量,
  • 是第张图像中错误预测为类别的像素数量,
  • 是第张图像中属于类别但被预测为其他类别的像素数量。

平均交并比()通过所有类别的IoU得分的平均值提供了一个单一的评价指标,定义如下:

其中是类别数量。

4 Architecture Discussion

在本节中,作者分析DDU-Net的一些架构特性。首先,作者通过实验和理论两方面来研究作者方法的内存需求,并将其与标准的U-Net模型进行比较。接着,作者分析具有不同架构的DDU-Net模型接收域的大小,即子网络和通信网络的深度不同时的情形。

Memory requirements

作者比较了图1所示的基础U-Net架构与所提出的DDU-Net模型架构的内存需求,两者均包含4个上采样和下采样块,操作于2个子图像;参见第3节及图2至图4。作者特别针对具有通信特征图的配置呈现结果。两个模型中的编码器-解码器网络的通道分布遵循图1所示的同一路线图。

表3详细分析了U-Net架构编码器和解码器在训练图像时的内存需求,这些编码器和解码器与DDU-Net中使用的相同。该表还包括了 Proposal 的具有的通信网络的内存需求。此表中的数值是基于理论推导(例如,见[51])并通过使用torch库[47]进行实验验证得出的。

分析表明,存储特征图的内存需求显著高于存储模型权重。浅层(包括输入块、第一个解码块和最后一个解码块)共同占据了特征图总内存分配近一半的比例。相比之下,由于U-Net更深层中拥有更多核和更大核尺寸,权重数量随之增加。在通信层使用的粗特征图导致通信网络相对较低的内存需求,证明了 Proposal 的DDU-Net在内存利用方面的效率。图7展示了在推理过程中,基础U-Net架构以及包含U-Net编码器-解码器和 Proposal 的通信网络的GPU在不同图像分辨率下的峰值内存使用情况。对于较小的图像分辨率,内存分配主要归因于模型权重。然而,当图像大小超过中等尺寸()像素时,存储特征图的内存需求开始主导总内存成本。在此之后,如图7所示,内存需求显著增加。从图7明显可以看出,随着分辨率的提高,DDU-Net与U-Net的内存需求相对差异变得较小:在32像素的分辨率下,相对差异为25.9%;而在4 096像素下,这一差异减少至仅0.65%。此外,值得注意的是,峰值内存随分辨率的增加呈二次方增长;因此,分辨率加倍将导致内存消耗增加四倍。

结果表明,尽管增加了通信网络的开销, Proposal 的DDU-Net仅需相对较小的内存额外负担,相较于基础U-Net而言。这展示了DDU-Net在内存利用上的高效性:它促进了子图像间的通信,同时仅轻微地增加了内存需求。

Receptive field analysis

所提出的通信网络在编码器-解码器网络的粗糙瓶颈层上运行。这些层的输出通常具有较大的感受野,因为它们是多次卷积和下采样操作的结果。因此,通信网络在增加模型架构的感受野大小方面非常有效。在表2中,作者根据[48]中提出的方法,比较了标准U-Net与所提出模型在不同深度的编码器-解码器网络中的(理论)感受野大小。对于DDU-Net,作者计算了一个无限大子域的理论感受野大小,为网络在处理多个子图像时真实感受野大小的上限提供了参考。比较结果显示,通信网络显著增大了编码器-解码器网络的感觉野大小。

5 Experimental Results

在本节中,作者根据第3节引入的数据集,比较了DDU-Net方法与相应的 Baseline U-Net模型在分割质量上的表现。此外,作者还进行了一项消融研究,以检验处理子图像的编码器-解码器网络之间通信的有效性。

Synthetic Dataset Results

在本节中,作者展示了使用第3.4.1小节所述的综合数据集所获得的结果。该数据集包括不同尺寸的灰度图像,具体为像素,其中。对于每个值,作者分别训练了 Baseline U-Nets和DDU-Nets。

在实验中,作者使用整个图像作为输入来训练 Baseline U-Nets。相比之下,DDU-Net架构接收个大小为像素的子图像,这些子图像由个独立的编码器-解码器克隆进行处理。U-Net和DDU-Net模型均设计用于接收1通道的灰度输入图像,并生成对应于像素类别_logit_的三通道输出:_背景_、_线段_和_圆_。

作者评估以下方面:

  • 在编码器-解码器网络中改变下采样和上采样块的数量(深度)的影响,
  • 通信网络深度的作用,这等同于发送到通信网络的特征图数量(表示为),
  • 特征图间通信的影响(,其中表示启用,表示禁用)。

如第3.2小节所述,作者使用表示法DDU-Net()来指代具有参数,,和的DDU-Net。对于深度为的 Baseline U-Net,作者使用表示法U-Net()。训练网络的超参数列于表7中。不同实验中编码器、解码器和通信网络的权重数量可以在附录中找到,具体在表5中。

5.1.1 Qualitative results

为了展示具有通信功能的粗网络的有效性,作者比较了三种架构在图8中对像素图像()的定性分割 Mask :U-Net(3)、DDU-Net()和DDU-Net()。作者观察到,由于缺乏通信,DDU-Net()无法正确预测线段的位置,因为它不知道另一子图像中的圆圈位置。相比之下,DDU-Net()正确预测了线段的位置。

图8显示,所提出的编码器-解码器网络能够捕获在通信特征图中的相关特征,并且通信网络能够在编码器-解码器网络之间传递这些信息。尽管在此示例中作者当前在DDU-Net()的瓶颈处通信所有32个特征图,但作者在5.1.2小节中将研究改变特征图数量所带来的影响。这一点尤其值得关注,因为通信仅在非常粗的 Level 上发生(使用像素的特征图),而线段是在更细粒度的分辨率上绘制的。显然,这种粗 Level 的通信对于解码器在细 Level 上产生良好的分割结果是足够的。

5.1.2 Varying the number of subimages and communicated feature maps

在图9中,展示了线段类别在不同子图像数量情况下, Baseline U-Net(3)以及DDU-Net()(其中特征图被通信)的IoU得分((3.2))。回想一下,的情况对应于通信网络深度为0的DDU-Net,意味着通信网络没有被使用,这也暗示编码器-解码器克隆之间没有通信。作者观察到了两个现象。首先,通信特征图的数量与结果质量之间存在明显的正相关关系。其次,对于较大的图像和子图像数量,DDU-Net的表现甚至优于 Baseline U-Net。

这与表2中报告的感受野大小有关:DDU-Net具有更大的感受野(像素对比像素),当圆圈相隔较远时,能获得更好的结果。

5.1.3 Impact of the depth of the encoder-decoder network

在DDU-Net中,编码器-解码器网络的深度为,即上采样和下采样块的数量。理想的深度是一种权衡:尽管浅层深度由于有限的感受野大小无法捕获所有相关特征(见表2),但过深的编码器-解码器包含更多参数,并可能产生过于粗糙且信息量较少的特征图以进行通信。

为了观察编码器-解码器深度对预测的影响,作者将DDU-Net()生成的分割 Mask 与U-Net()生成的分割 Mask 进行比较,其中。这些预测的结果展示在图10中。对于具有的编码器-解码器网络,作者观察到线段仅在位于两个圆圈之间中心区域的像素上预测正确。这也可以通过感受野来解释:浅层网络具有有限的感受野,以至于远离两个圆圈之一的像素不受该圆圈的影响;这导致了错误的分割结果。也可以看出,启用通信的DDU-Net比和时的U-Net得到更好的结果,因为通信网络增加了感受野的大小。

5.1.4 Generalization to different numbers of subimages

Baseline U-Net和DDU-Net都是尺寸无关的,因为它们都是全卷积神经网络。这使得它们能够处理与训练图像尺寸不同的图像。在DDU-Net的情况下,这意味着作者可以改变子图像的大小和子图像的数量。为了测试DDU-Net关于不同通信量下子图像数量的泛化能力,作者针对每个,按照3.4.1节所述,在空间尺寸为像素的输入图像上训练一个DDU-Net,其中。然后,每个模型都在包含6个子图像的数据集上进行评估。结果展示在图11中。

作者可以得出以下观察:

  1. 在3、4和6个子图像上训练的模型在6个子图像的测试数据集上表现最好。这可能是由于这些数据集中的图像足够小,可以被DDU-Net()的像素感受野几乎完全覆盖;参见观察2。
  2. 对于8和16个子图像的较不准确结果可以通过DDU-Net()的有限感受野大小像素来解释;参见表2。
  3. 当模型仅在两个子图像上进行训练,然后在对更多子图像进行评估时,IoU得分会下降。这可能是由于训练数据集中的每个子图像都包含一个圆和一个线段的片段;然而,没有包含不含圆或线段的子图像的样本。

这些发现表明,DDU-Net可以在与评估时不同的子图像数量上进行训练。

5.1.5 Summary of results on the synthetic dataset

以下是对合成数据集的研究结果的总结:

  • 通信网络能够跨子图像传递上下文信息,这一点从图8的定性比较和图9的定量结果中都变得清晰。
  • 当参与通信的特征图数量增加时,分割质量也随之提高。
  • 相比于 Baseline U-Net,具有通信功能的DDU-Net更大的感受野尺寸导致了更优的结果;参见图9和图10。
  • DDU-Net可以在固定数量的子图像上进行训练,并在不同数量的子图像上进行评估。

DeepGlobe Dataset

现在,作者评估DDU-Net在处理具有高分辨率图像的真实世界数据集上的有效性:DeepGlobe数据集;参见第3.4.2小节。

鉴于DeepGlobe数据集相对于分割任务复杂性的803张图像数量较少,作者对编码器-解码器架构进行了几项调整,包括批量归一化、随机丢弃层和数据增强(水平翻转、垂直翻转、随机旋转弧度,其中),并且作者采用了一个预训练的图像编码器模型,即ResNet-18 [24]。这个模型在ImageNet数据集[14]上针对超过1000个类别进行了训练,为各种图像提供了一个丰富的特征表示。ResNet-18模型由带有跳过连接的残差块组成,允许有效的梯度流动。作者使用ResNet-18的前四个残差块来初始化作者的编码器。这种策略利用了预训练模型中的预存知识,以增强网络从有限数据集中泛化模式的能力。图12给出了所采用模型架构的可视化。关于网络每个组件中模型参数的详细分布概览。

在训练期间,作者保持了ResNet-18模型的权重不变,只调整了解码器和通信网络的权重。此外,作者在瓶颈层中插入了两个额外的卷积层。这些层允许网络重新构造和精炼由ResNet编码器产生的特征图,为通信网络生成相关信息。

为了确保 Baseline U-Net与DDU-Net的分割质量之间的公平比较,作者都在大小相等的图像上训练了这两种模型。由于 Baseline U-Net架构本质上无法在多个GPU之间并行化,且训练设备内存有限(32 GB),作者从DeepGlobe数据集中裁剪出的非重叠块,从而得到包含张图像的训练数据集。区分DDU-Net架构中的这些大型“全局块”和子图像至关重要。在DDU-Net的训练过程中,全局块进一步划分为较小的子图像,然后它们被分配到编码器-解码器网络中。采用混合精度训练[43],这种方法允许作者在单个GPU上以12张图像的小批量进行训练。其他用于训练的超参数显示在表4中。此外,考虑到网络参数的随机初始化,作者用相同的设置和训练数据集重复训练每个网络三次。对于每种配置,选择在测试数据集上IoU得分表现最佳的模型。

作者想要调查DDU-Net是否能够与 Baseline U-Net表现相当。然而,当作者在DDU-Net架构中包含通信时,额外的可训练参数意味着潜在的学习更复杂数据模式的能力。为了隔离(1)通信网络中“额外参数”的影响和(2)“通信本身”的影响,作者改变了DDU-Net和DDU-Net中的特征图数量。DDU-Net场景表示带有瓶颈层额外粗网络的编码器-解码器网络(粗网络选择与通信网络相同),但 子域之间 没有 信息交换。在这种情况下,粗网络仅操作于 局部 瓶颈特征图,而不是所有子图像的拼接特征图。值得注意的是,对应于 Baseline U-Net。启用通信()和禁用通信()的情况之间的差异显示了跨子图像通信的影响。

5.2.1 Quantitative results

图13展示了在不同DDU-Net配置和图像划分下,在图像上训练并在固定子图像大小的测试数据集上评估的平均IoU得分,所采用的方法是在5.1.4小节中讨论的。

从这幅图中,作者可以得出两个主要观察结果。首先,随着粗网络中特征图数量的增加,分割质量得到提升。这是预料之中的,因为特征图数量的增加会导致参数数量的增加和感受野的扩大,使模型能够捕捉到更复杂和更远的模式。图13中的另一个趋势是,随着子图像数量的增加,或者等价地说,随着子图像大小的减小,预测质量会下降。然而,对于DDU-Net,与DDU-Net相比,质量下降要小得多。这表明,结合粗网络的通信有效地在子图像之间传递了上下文信息,对于这个真实数据集也是如此。

5.2.2 Qualitative results

在图14中,展示了不同训练和模型配置下的预测示例。作者观察到,在上训练的DDU-Net以及启用了通信的子图像上训练的DDU-Net都能产生较好的分割结果,尽管与真实 Mask 存在一些差异。相反,没有通信的子图像上训练的DDU-Net显示出较差的预测效果。为了更清晰地说明预测误差,作者在附录中包含了图15,其中黑色表示正确预测,白色表示错误预测。关键区别在于相邻子图像之间的一致性:当启用通信时,相邻子图像的预测边界更为平滑,而不是呈现块状模式。这突显了在DDU-Net架构中子图像之间通信的有效性。

6 Conclusion

本文提出了一种基于域分解的U-Net(DDU-Net)架构,用于语义分割任务。

作者的结果显示,通过包含子图像间的通信,DDU-Net可以高效地处理高分辨率图像分割,而不会牺牲准确度或内存效率。作者的方法通过利用子图像间的通信来提高分割准确度。

未来的研究将聚焦于改进通信策略,将DDU-Net应用于更复杂的数据库,并评估并行化对计算时间的好处。另一个研究方向是将本文中开发的并行化策略扩展到其他编码器-解码器架构。

参考

[1].DDU-Net: A domain decomposition-based CNN for high-resolution image segmentation on multiple GPUs.

0 人点赞