使用PyTorch进行高效图像分割:第一部分

概念和思想

在这个由4个部分组成的系列中,我们将使用PyTorch中的深度学习技术,逐步地从头开始实现图像分割。本文将从基本概念和图像分割所需的思想开始这个系列。

Figure 1: Pet images and their segmentation masks (Source: The Oxford-IIIT Pet Dataset )

与Naresh Singh共同撰写

图像分割是一种将属于图像中特定物体的像素孤立出来的技术。孤立出物体像素可以打开有趣的应用程序的大门。例如,图1中右侧的图像是对应于左侧宠物图像的蒙版,在这些蒙版中,黄色像素属于宠物。一旦识别出像素,我们就可以轻松地使宠物变大或更改图像背景。这种技术在几个社交媒体应用程序的面部滤镜功能中广泛使用。

本系列文章的目标是让读者了解构建视觉AI模型并使用PyTorch进行不同设置的实验所需的所有步骤。

本系列文章

本系列文章适合所有深度学习经验水平的读者。如果您想学习关于深度学习和视觉AI的实践以及一些扎实的理论和实践经验,那么您来对地方了!预计这将是一个由以下文章组成的4部分系列:

  1. 概念和思想(本文)
  2. 基于CNN的模型
  3. 深度可分离卷积
  4. 基于Vision Transformer的模型

图像分割简介

图像分割将图像分割为对应于物体、背景和边界的区域。请看图2,它显示了一幅城市场景。它使用不同颜色的蒙版标记了对应于汽车、摩托车、树木、建筑物、人行道和其他有趣物体的区域。这些区域是通过图像分割技术识别出来的。

历史上,我们使用专门的图像处理工具和流程将图像分解成区域。然而,由于过去二十年中视觉数据的惊人增长,深度学习已经成为解决图像分割任务的首选解决方案。它显着减少了依赖专家构建特定领域图像分割策略的依赖,这是过去所做的。如果有足够的训练数据可以用于任务,深度学习实践者可以训练图像分割模型。

Figure 2: A segmented scene from the a2d2 dataset (CC BY-ND 4.0)

图像分割的应用

图像分割在通信、农业、交通、医疗保健等多个领域都有应用。此外,随着视觉数据的增长,其应用也在不断增加。以下是一些例子:

  • 自动驾驶汽车中,深度学习模型不断处理汽车摄像头的视频流,将场景分割成汽车、行人和交通信号等对象,这对汽车的安全运行至关重要。
  • 医学成像中,图像分割帮助医生识别与肿瘤、病变和其他异常对应的医学扫描区域。
  • Zoom视频通话中,它用于通过虚拟场景替换背景,以保护个人隐私。
  • 农业中,使用图像分割识别杂草和作物区域的信息用于维护健康的农作物产量。

您可以在v7labs的这个页面上阅读更多关于图像分割的实际应用的细节。

图像分割任务有哪些不同类型?

有许多不同类型的图像分割任务,每种任务都有其优缺点。最常见的两种图像分割任务是:

  • 类别或语义分割:类别分割为每个图像像素分配一个语义类别,如背景、道路、汽车或人。如果图像中有两辆车,则与两辆车对应的像素将被标记为汽车像素。它通常用于自动驾驶和场景理解等任务。
  • 对象或实例分割:对象分割识别对象并为图像中的每个唯一对象分配一个掩码。如果图像中有两辆车,则与每辆车对应的像素将被标识为属于不同对象。对象分割通常用于跟踪单个对象,例如自动驾驶汽车编程以跟踪其前方的特定车辆。
Figure 3: Object and class segmentations (Source: MS Coco — Creative Commons Attribution license )

在本系列中,我们将专注于类别分割。

实现高效图像分割所需的决策

为了实现速度和准确性的高效训练,需要在项目的生命周期中做出许多重要决策,包括(但不限于):

  1. 选择深度学习框架
  2. 选择良好的模型架构
  3. 选择优化所关心方面的有效损失函数
  4. 避免过拟合和欠拟合
  5. 评估模型的准确性

在本文的其余部分,我们将深入探讨上述每个方面,并提供许多链接,以便更详细地讨论每个主题。

PyTorch 用于高效图像分割

什么是 PyTorch?

“PyTorch 是一个开源的深度学习框架,旨在为研究提供灵活和模块化的结构,具有生产部署所需的稳定性和支持。PyTorch 提供了一个 Python 包,用于高级特性,如张量计算(类似于 NumPy)和强大的 GPU 加速以及 TorchScript,可在急切模式和图形模式之间轻松转换。随着 PyTorch 的最新发布,该框架提供了基于图形的执行、分布式训练、移动部署和量化。”(出自 Meta AI 页面关于 PyTorch 的介绍)

PyTorch 是用 Python 和 C++ 编写的,易于使用和学习,同时运行效率也很高。它支持多种硬件平台,包括(服务器和移动)CPU、GPU 和 TPU。

为什么 PyTorch 是进行图像分割的好选择?

PyTorch 是进行深度学习研究和开发的流行选择,因为它提供了一个灵活和强大的环境来创建和训练神经网络。对于基于深度学习的图像分割,它是一个极佳的框架选择,原因如下:

  • 灵活性:PyTorch 是一个灵活的框架,可以通过多种方式创建和训练神经网络。您可以使用预训练模型,也可以轻松地从头开始创建自己的模型。
  • 后端支持:PyTorch 支持多个后端,如 GPU/TPU 硬件。
  • 领域库:PyTorch 有丰富的领域库,使得处理特定数据垂直领域非常容易。例如,对于与视觉(图像/视频)相关的 AI,PyTorch 提供了一个名为 torchvision 的库,在本系列中我们将广泛使用。
  • 易用性和社区采用率:PyTorch 是一个易于使用的框架,有着良好的文档和庞大的用户和开发者社区。许多研究人员使用 PyTorch 进行实验,并且他们的发表论文中的模型实现在 PyTorch 中是免费提供的。

数据集的选择

我们将使用牛津IIIT宠物数据集(根据CC BY-SA 4.0许可)进行类分割。该数据集共有3680张训练集图像,每个图像都有一个分割trimap与其关联。trimap是三个像素类之一:

  1. 宠物
  2. 背景
  3. 边框

我们选择此数据集,因为它足够多样化,可以提供我们一个非平凡的类分割任务。此外,它不是太复杂,以至于我们花费时间处理类不平衡等问题,而忘记了我们想要学习和解决的主要问题,即类分割。

用于图像分割任务的其他流行数据集包括:

  1. Pascal VOC(视觉对象类)
  2. MS Coco
  3. Cityscapes

使用PyTorch高效进行图像分割

在本系列中,我们将从头开始训练多个用于类分割的模型。构建和训练模型时需要考虑许多因素。下面,我们将查看一些在这样做时需要做出的关键决策。

选择适合您任务的正确模型

选择用于图像分割的正确深度学习模型时需要考虑许多因素。一些最重要的因素包括:

  • 图像分割任务的类型:图像分割任务有两种主要类型:类(语义)分割和对象(实例)分割。由于我们专注于更简单的类分割问题,因此我们将考虑相应地建模我们的问题。
  • 数据集的大小和复杂性:数据集的大小和复杂性将影响我们需要使用的模型的复杂性。例如,如果我们使用具有较小空间维度的图像,则可以使用更简单(或更浅)的模型,例如全卷积网络(FCN)。如果我们使用一个大型且复杂的数据集,则可以使用更复杂(或更深的)模型,例如U-Net。
  • 可用的预训练模型:有许多可用于图像分割的预训练模型。这些模型可以用作我们自己模型的起点,或者我们可以直接使用它们。但是,如果我们使用预训练模型,则可能会受到模型输入图像的空间尺寸的限制。在本系列中,我们将专注于从头开始训练模型。
  • 可用的计算资源:深度学习模型的训练可能需要大量计算资源。如果我们计算资源有限,则可能需要选择更简单的模型或更高效的模型架构。

在本系列中,我们将使用牛津IIIT宠物数据集,因为它足够大,可以让我们训练一个小型模型,并需要使用GPU。我们强烈建议在kaggle.com上创建帐户,或者使用Google Colab的免费GPU来运行本系列中引用的笔记本和代码。

模型架构

以下是用于图像分割的最流行的深度学习模型架构:

  • U-Net:U-Net是常用于图像分割任务的卷积神经网络。它使用跳跃连接,可以帮助更快地训练网络并获得更好的整体准确性。如果必须选择,U-Net始终是一个出色的默认选择
  • FCN:全卷积网络(FCN)是一个完全卷积网络,但不像U-Net那样深。缺乏深度主要是因为在更高的网络深度时,准确性会下降。这使得它更快速进行训练,但可能不像U-Net那样准确。
  • SegNet:SegNet是一种流行的模型架构,类似于U-Net,并且使用的激活内存比U-Net少。我们将在本系列中使用SegNet。
  • Vision Transformer(ViT):由于其简单的结构和关注机制在文本、视觉和其他领域的适用性,视觉变换器最近变得越来越流行。与卷积神经网络相比,视觉变换器在训练和推断方面可能更高效,但历史上需要更多的数据来训练。我们也将在本系列中使用ViT。
图4:U-Net模型架构。来源:弗莱堡大学,U-Net的原始作者。

这些仅是可用于图像分割的许多深度学习模型中的一小部分。您特定任务的最佳模型将取决于早期提到的因素、特定任务和您自己的实验。

选择正确的损失函数

对于图像分割任务,损失函数的选择是重要的,因为它可能对模型的性能产生显着的影响。有许多不同的损失函数可用,每种函数都有其优点和缺点。用于图像分割的最流行的损失函数包括:

  • 交叉熵损失:交叉熵损失是预测概率分布与真实概率分布之间差异的度量
  • IoU损失:IoU损失测量预测掩模和每类别的真实掩模之间的重叠量。IoU损失会惩罚预测或召回任何一方的情况。定义的IoU不可微分,因此我们需要稍微调整它以将其用作损失函数
  • Dice损失:Dice损失也是预测掩模和真实掩模之间重叠的度量。
  • Tversky损失:Tversky损失被提出作为一种稳健的损失函数,可用于处理不平衡的数据集。
  • Focal损失:Focal损失旨在专注于难以分类的样本。这对于改善模型在具有挑战性的数据集上的性能可能有所帮助。

特定任务的最佳损失函数将取决于任务的特定要求。例如,如果准确性更重要,则IoU损失或Dice损失可能是更好的选择。如果任务不平衡,则Tversky损失或Focal损失可能是不错的选择。所使用的特定损失函数可能会影响在训练模型时收敛速率。

损失函数是模型的超参数,根据我们看到的结果使用不同的损失可以让我们更快地减少损失并提高模型的准确性。

默认:在这个系列中,我们将使用交叉熵损失,因为当结果未知时,这通常是一个好的默认选择。

您可以使用以下资源了解更多有关损失函数的信息。

  1. PyTorch损失函数:终极指南
  2. Torchvision – 损失
  3. Torchmetrics

让我们详细了解下面定义的IoU损失,作为分割任务的强大替代方案。

自定义IoU损失

IoU定义为交集除以并集。对于图像分割任务,我们可以通过计算(对于每个类别)由模型预测的该类别中的像素与地面实况分割掩模中的像素的交集来计算此值。

例如,如果我们有2个类别:

  1. 背景
  2. 人物

然后,我们可以确定哪些像素被分类为人物,并将其与人物的实际像素进行比较,计算人物类别的IoU。同样,我们可以计算背景类别的IoU。

一旦我们拥有这些特定类别的IoU指标,我们可以选择在加权之前对它们进行未加权平均,以考虑任何类型的类别不平衡,就像我们在之前的示例中看到的那样。

所定义的IoU度量需要我们为每个度量计算硬标签。这需要使用argmax()函数,该函数不可微分,因此我们无法将此度量用作损失函数。因此,我们不使用硬标签,而是应用softmax()并使用预测概率作为软标签来计算IoU度量。这会产生可微分的度量,我们可以从中计算损失。因此,有时,当在损失函数的上下文中使用时,IoU度量也称为软IoU度量。

如果我们有一个介于0.0和1.0之间的指标(M),我们可以计算损失(L):

L = 1 — M

但是,如果您的指标的值介于0.0到1.0之间,还有另一种转换指标为损失的技巧。计算:

L = -log(M)

即计算指标的负对数。这与先前的公式有实质性区别,您可以在这里和这里阅读更多信息。基本上,它会使您的模型学习更好。

Figure 6: Comparing the loss resulting from 1-P(x) with -log(P(x)). Source: Author(s).

使用IoU作为损失函数也使损失函数更接近于捕捉我们真正关心的内容。使用评估指标作为损失函数有利有弊。如果您有兴趣进一步探索此领域,可以从此处的stackexchange讨论开始。

数据增强

为了以高精度有效地训练模型,需要注意用于训练模型的训练数据的数量和类型。所使用的训练数据的选择将极大地影响最终模型的准确性,因此,如果您想从本系列文章中学到一件事,那应该是这个!

通常,我们将数据分成三部分,部分的比例大致如下。

  1. 训练(80%)
  2. 验证(10%)
  3. 测试(10%)

您将在训练集上训练模型,在验证集上评估准确性,并重复此过程,直到您对报告的指标感到满意。然后,您才会在测试集上评估模型,然后报告数字。这样做是为了防止任何偏见渗入您的模型体系结构和在培训和评估期间使用的超参数。一般来说,您越是根据测试数据中看到的结果调整设置,您的结果就会变得越不可靠。因此,我们必须将决策限制在我们在训练和验证数据集上看到的结果。

在本系列中,我们将不使用测试数据集。相反,我们将使用我们的测试数据集作为验证数据集,并对测试数据集应用数据增强,以便我们始终在略有不同的数据上验证我们的模型。这有点像一个技巧,我们这样做只是为了快速迭代和捷径。对于生产模型开发,您应该尽量遵循上述标准配方。

在本系列中,我们将使用的数据集中有3680个图像在训练集中。虽然这可能看起来像是一大批图像,但我们希望确保我们的模型不会在这些图像上过度拟合,因为我们将对模型进行多次训练。

在单个训练时期内,我们会对整个训练数据集上的模型进行训练,而我们通常会在生产中对模型进行60次或更多次训练。在本系列中,我们只对模型进行20次训练以加快迭代时间。为了防止过度拟合,我们将使用一种称为数据增强的技术,该技术用于从现有输入数据生成新的输入数据。图像输入的数据增强的基本思想是,如果您稍微改变图像,那么对于模型来说,它感觉像一个新图像,但您可以推断预期的输出是否相同。以下是我们将在本系列中应用的一些数据增强示例。

  1. 随机水平翻转
  2. 随机颜色抖动

虽然我们将使用Torchvision库应用上述数据增强,但我们鼓励您在视觉任务中评估Albumentations数据增强库。两个库都有丰富的可用于图像数据的转换集。我们个人继续使用Torchvision,只是因为它是我们开始使用的。Albumentations支持更丰富的数据增强原语,可以同时更改输入图像以及基本事实标签或掩码。例如,如果您要调整图像大小或翻转图像,则需要同时对基本事实分割掩码进行相同的更改。Albumentations可以为您完成这项工作。

总的来说,这两个库都支持对图像进行像素级转换或更改图像的空间维度的转换。Torchvision将像素级变换称为颜色变换,将空间变换称为几何变换。

下面,我们将看到Torchvision和Albumentations库应用的一些像素级和几何变换的示例。

图7:使用Albumentations应用像素级数据增强的示例。来源:Albumentations
图8:使用Torchvision变换应用于图像的数据增强示例。来源:作者(笔记本)
图9:使用Albumentations应用的空间级别变换的示例。来源:作者(笔记本)

评估模型性能

在评估模型性能时,您需要知道模型在代表真实世界数据的质量度量上的表现如何。例如,对于图像分割任务,我们想要知道模型能够准确预测每个像素的正确类别的准确度。因此,我们称像素准确度为此模型的验证度量。

您可以使用评估度量作为损失函数(为什么不优化您真正关心的内容!),但这可能并不总是可能的。

除了准确度之外,我们还将跟踪IoU度量(也称为Jaccard指数)和我们上面定义的自定义IoU度量。

有关适用于图像分割任务的各种精度度量的更多信息,请参见:

  • 所有分割度量-Kaggle
  • 如何评估图像分割模型
  • 评估图像分割模型

将像素准确度作为性能度量的缺点

虽然准确度度量可能是衡量图像分割任务性能的好默认选择,但它确实有其自身的缺点,根据您的特定情况可能非常重要。

例如,考虑一个图像分割任务,以识别图片中人眼,并相应地标记那些像素。因此,模型将将每个像素分类为以下之一:

  1. 背景
  2. 眼睛

假设每个图像中只有一个人,并且98%的像素不对应于眼睛。在这种情况下,模型可以简单地学习预测每个像素作为背景像素,并在分割任务上实现98%的像素准确度。哇!

图10:一个人脸的图像和对应的眼睛分割掩码。您可以看到眼睛占整个图像的一小部分。来源:改编自Unsplash

在这种情况下,使用IoU或Dice指标可能是一个更好的主意,因为IoU将捕获多少预测是正确的,并且不会受到每个类别或类别在原始图像中占据的区域的偏见。您甚至可以考虑使用每个类别的IoU或Dice系数作为度量。这可能更好地捕捉您的模型在手头任务上的表现。

当仅考虑像素精度时,我们要计算分割掩模的目标对象(如上面的眼睛)的精确度和召回率可以捕捉到我们正在寻找的细节。

现在我们已经涵盖了图像分割的大部分理论基础,让我们转入与推理和部署图像分割相关的考虑。

模型大小和推理延迟

最后,我们要确保我们的模型具有合理数量的参数,但不要太多,因为我们想要一个小而高效的模型。我们将在将来的文章中更详细地讨论使用高效模型架构来减小模型大小的方面。

就推理延迟而言,我们模型执行的数学操作(mult-adds)是重要的。可以使用torchinfo包来显示模型大小和mult-adds。虽然mult-adds是确定模型延迟的良好代理,但在各种后端上的延迟可能有很大的变化。在特定的设备上使用预期在生产环境中看到的输入集对其进行分析和基准测试是确定模型在特定后端或设备上的性能的唯一真正方法。

from torchinfo import summarymodel = nn.Linear(1000, 500)summary(  model,  input_size=(1, 1000),  col_names=["kernel_size", "output_size", "num_params", "mult_adds"],  col_width=15,)

输出:

====================================================================================================Layer (type:depth-idx)                   Kernel Shape    Output Shape    Param #         Mult-Adds====================================================================================================Linear                                   --              [1, 500]        500,500         500,500====================================================================================================Total params: 500,500Trainable params: 500,500Non-trainable params: 0Total mult-adds (M): 0.50====================================================================================================Input size (MB): 0.00Forward/backward pass size (MB): 0.00Params size (MB): 2.00Estimated Total Size (MB): 2.01====================================================================================================

更多阅读

以下文章提供有关图像分割基础的其他信息。如果您是喜欢阅读相同主题的不同观点的人,请考虑阅读它们。

  1. 计算机视觉中的图像分割指南:最佳实践
  2. 图像分割简介:深度学习 vs. 传统[+示例]
  3. 图像分割:基础知识和5个关键技术

如果您希望亲自动手使用Oxford IIIT宠物数据集并使用torchvision和Albumentations执行图像增强,我们在Kaggle上提供了一个入门笔记本,您可以克隆并尝试使用。本文中的许多图像都是由该笔记本生成的!

文章概述

以下是我们迄今为止讨论的内容的快速概述。

  • 图像分割是一种将图像分割为多个部分的技术(来源:维基百科)
  • 图像分割任务主要有两种类型:类(语义)分割和对象(实例)分割。类分割将图像中的每个像素分配到一个语义类别中。对象分割识别图像中的每个独特对象并为每个独特对象分配一个掩模
  • 我们将在本系列高效图像分割中使用PyTorch作为深度学习框架和Oxford IIIT宠物数据集
  • 选择适合图像分割的正确深度学习模型时,需要考虑许多因素,包括(但不限于)图像分割任务的类型,数据集的大小和复杂性,预训练模型的可用性以及计算资源。一些最流行的图像分割深度学习模型架构包括U-Net、FCN、SegNet和Vision Transformer(ViT)
  • 对于图像分割任务,选择损失函数是重要的,因为它可以对模型的性能和训练效率产生重大影响。对于图像分割任务,我们可以使用交叉熵损失、IoU损失、Dice损失或Focal损失(等等)
  • 数据增强是一种有价值的技术,用于防止过拟合以及处理不足的训练数据
  • 评估模型的性能对于当前任务非常重要,必须仔细选择此指标
  • 模型大小和推理延迟是开发模型时需要考虑的重要指标,特别是如果您打算将其用于实时应用程序,例如面部分割或背景噪声去除

在下一篇文章中,我们将使用PyTorch从头构建一个卷积神经网络(CNN),在Oxford IIIT宠物数据集上执行图像分割。