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

深度可分离卷积

在这个由四部分组成的系列中,我们将使用PyTorch中的深度学习技术从头开始逐步实现图像分割。本篇将着重优化我们的CNN基线模型,使用深度可分离卷积来减少可训练参数的数量,使模型可以部署在手机和其他边缘设备上。

与Naresh Singh共同合著

Figure 1: Result of running image segmentation using a CNN with depth-wise separable convolutions instead of regular convolutions. From top to bottom, input images, ground truth segmentation masks, and predicted segmentation masks. Source: Author(s)

文章大纲

在本文中,我们将增强我们之前构建的卷积神经网络(CNN),以减少网络中可学习参数的数量。识别输入图像中的宠物像素(属于猫、狗、仓鼠等的像素)的任务保持不变。我们选择的网络仍然是SegNet,我们唯一的改变是用深度可分离卷积(DSC)替换我们的卷积层。在这之前,我们将深入研究深度可分离卷积的理论和实践,欣赏该技术背后的思想。

在本文中,我们将参考此笔记本中的代码和结果进行模型训练,以及此笔记本中的DSC入门。如果您想要重现结果,您需要一台GPU确保第一个笔记本在合理的时间内完成运行。第二个笔记本可以在普通CPU上运行。

系列文章

本系列文章适合所有深度学习的阅读者。如果您想一边学习深度学习和视觉人工智能,一边了解一些扎实的理论和实践经验,您来对了地方!本系列将包括以下四篇文章:

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

介绍

让我们从模型大小和计算成本的角度更加深入地了解卷积。可训练参数的数量是模型大小的一个很好的指标,而张量操作的数量反映了模型的复杂度或计算成本。假设我们有一个卷积层,其中包含n个尺寸为dₖ x dₖ的滤波器。进一步假设该层处理形状为m x h x w的输入,其中m是输入通道的数量,而h和w分别是高度和宽度维度。在这种情况下,卷积层将产生一个形状为n x h x w的输出,如图2所示。我们假设卷积使用stride=1。让我们继续评估这个设置的可训练参数和计算成本。

Figure 2: Regular convolutional filters applied to input to produce output. Assume stride=1 and padding=dₖ-2. Source: Efficient Deep Learning Book

可训练参数的评估:我们有n个滤波器,每个滤波器都有m x dₖ x dₖ个可学习参数。这将产生n x m x dₖ x dₖ个可学习参数。为简化此讨论,我们忽略了偏差项。让我们看一下下面的PyTorch代码,以验证我们的理解。

import torchfrom torch import nndef num_parameters(m):return sum([p.numel() for p in m.parameters()])dk, m, n = 3, 16, 32print(f"Expected number of parameters: {m * dk * dk * n}")conv1 = nn.Conv2d(in_channels=m, out_channels=n, kernel_size=dk, bias=False)print(f"Actual number of parameters: {num_parameters(conv1)}")

打印以下内容。

期望参数数量:4608实际参数数量:4608

现在,让我们评估卷积的计算成本。

计算成本的评估:当在输入大小为 h x w 的输入上以 stride=1 和 padding=dₖ-2 运行形状为 m x dₖ x dₖ 的单个卷积滤波器时,每个大小为 dₖ x dₖ 的图像部分将应用卷积滤波器 h x w 次,总计 h x w 个部分。每个滤波器或输出通道的成本为 m x dₖ x dₖ x h x w,因此希望计算 n 个输出通道,总成本将为 m x dₖ x dₖ x h x n。让我们使用 torchinfo PyTorch 包验证这一点。

from torchinfo import summaryh, w = 128, 128print(f"期望总乘积:{m * dk * dk * h * w * n}")summary(conv1, input_size=(1, m, h, w))

将打印以下内容。

期望总乘积:75497472==========================================================================================Layer (type:depth-idx)                   Output Shape              Param #==========================================================================================Conv2d                                   [1, 32, 128, 128]         4,608==========================================================================================Total params: 4,608Trainable params: 4,608Non-trainable params: 0Total mult-adds (M): 75.50==========================================================================================Input size (MB): 1.05Forward/backward pass size (MB): 4.19Params size (MB): 0.02Estimated Total Size (MB): 5.26==========================================================================================

如果我们暂时忽略卷积层的实现细节,我们会意识到,在高层次上,卷积层只是将 m x h x w 输入转换为 n x h x w 输出。通过逐步学习功能的可训练滤波器实现转换。随之而来的问题是:是否可能使用更少的可学习参数来实现这种转换,并同时确保在层的学习能力方面最小的妥协?深度可分离卷积被提出来回答这个问题。让我们详细了解它们,并了解它们在我们的评估指标中的表现。

深度可分离卷积

深度可分离卷积(DSC)的概念最初由 Laurent Sifre 在他们的博士论文中提出的,题为 Rigid-Motion Scattering For Image Classification。自那以后,它们已经成功地用于各种流行的深度卷积网络,如 XceptionNet 和 MobileNet。

常规卷积和 DSC 之间的主要区别在于,DSC 由以下 2 个卷积组成:

  1. 一个深度分组卷积,其中输入通道数 m 等于输出通道数,以使每个输出通道仅受单个输入通道的影响。在 PyTorch 中,这称为“分组”卷积。您可以在此处阅读有关 PyTorch 中分组卷积的更多信息。
  2. 一个点卷积(过滤器大小=1),其操作类似于常规卷积,以便每个 n 个滤波器在所有 m 个输入通道上运行,以产生单个输出值。
图3:应用深度可分离卷积滤波器以产生输出的输入。假设 stride=1 和 padding=dₖ-2。来源:高效的深度学习书

让我们为 DSC 执行与常规卷积相同的练习,计算可训练参数和计算数量。

可训练参数的评估:“分组”卷积具有 m 个滤波器,每个滤波器具有 dₖ x dₖ 可学习参数,可产生 m 个输出通道。这导致总共有 m x dₖ x dₖ 个可学习参数。点卷积具有 n 个大小为 m x 1 x 1 的滤波器,总量为 n x m x 1 x 1 可学习参数。让我们查看下面的 PyTorch 代码,以验证我们的理解。

class DepthwiseSeparableConv(nn.Sequential):    def __init__(self, chin, chout, dk):        super().__init__(            # 深度卷积            nn.Conv2d(chin, chin, kernel_size=dk, stride=1, padding=dk-2, bias=False, groups=chin),            # 逐点卷积            nn.Conv2d(chin, chout, kernel_size=1, bias=False),        )conv2 = DepthwiseSeparableConv(chin=m, chout=n, dk=dk)print(f"预期参数数量: {m * dk * dk + m * 1 * 1 * n}")print(f"实际参数数量: {num_parameters(conv2)}")

将会输出:

预期参数数量: 656实际参数数量: 656

可以看出,DSC版本的参数数量大约减少了7倍。接下来,让我们关注一下DSC层的计算成本。

计算成本的评估:假设输入的空间尺寸为m x h x w。在DSC的分组卷积段中,我们有m个大小为dₖ x dₖ的滤波器。每个滤波器应用于其相对应的输入通道,导致段的成本为 m x dₖ x dₖ x h x w。对于逐点卷积,我们应用了大小为 m x 1 x 1 的n个滤波器,以生成n个输出通道。这导致段的成本为 n x m x 1 x 1 x h x w。我们需要将分组和逐点运算的成本相加,以计算总成本。接下来,我们使用torchinfo的PyTorch软件包进行验证。

print(f"预期总乘法次数: {m * dk * dk * h * w + m * 1 * 1 * h * w * n}")s2 = summary(conv2, input_size=(1, m, h, w))print(f"实际乘法次数: {s2.total_mult_adds}")print(s2)

将会输出:

预期总乘法次数: 10747904实际乘法次数: 10747904==========================================================================================Layer (type:depth-idx)                   Output Shape              Param #==========================================================================================DepthwiseSeparableConv                   [1, 32, 128, 128]         --├─Conv2d: 1-1                            [1, 16, 128, 128]         144├─Conv2d: 1-2                            [1, 32, 128, 128]         512==========================================================================================总参数数: 656可训练参数数: 656不可训练参数数: 0总乘加次数 (M): 10.75==========================================================================================输入尺寸 (MB): 1.05前向/后向传递尺寸 (MB): 6.29参数尺寸 (MB): 0.00预计总尺寸 (MB): 7.34==========================================================================================

让我们比较一下几个示例的卷积大小和成本,以获得一些直觉。

Regular和Depthwise Separable Convolutions的大小和成本比较

为了比较正常卷积和深度可分离卷积的大小和成本,我们假设将大小为128 x 128的输入送入网络,内核大小为3 x 3,网络在逐渐减半空间尺寸并将通道维度加倍。我们假设每个步骤仅有一个2d-conv层,但实际上可能有更多。

Figure 4: Comparing the number of trainable parameters (size) and multi-adds (cost) of regular and depthwise separable convolutions. We also show the ratio of the size and cost for the 2 types of convolutions. Source: Author(s).

您可以看到,对于上述配置,DSC的大小和计算成本平均约为常规卷积的11%至12%。

图5:正常卷积与深度可分离卷积的相对大小和成本。来源:作者。

现在我们已经对各种卷积类型及其成本有了很好的理解,你可能会想知道使用深度可分离卷积是否有任何缺点。到目前为止,我们所见的一切似乎都表明它们在各个方面都更好!但是,我们还没有考虑到一个重要的方面,即它们对模型准确性的影响。让我们通过下面的实验来深入了解。

使用深度可分离卷积的 SegNet

本笔记本包含本部分的所有代码。

我们将从上一篇文章中适应我们的 SegNet 模型,并将所有常规卷积层替换为 DSC 层。一旦我们这样做,我们就会注意到我们笔记本中的参数数量从 15.27M 减少到 1.75M 参数,这是一个 88.5% 的减少!这符合我们先前估计的网络可训练参数减少 11% 到 12% 的估计结果。

在模型训练和验证期间使用与之前类似的配置。配置如下所示。

  1. 在训练集上应用随机水平翻转和颜色抖动数据增强,以防止过拟合
  2. 将图像调整大小为 128×128 像素,进行非保持纵横比的调整操作
  3. 不对图像应用输入规范化,而是使用批量规范化层作为模型的第一层
  4. 使用 Adam 优化器进行 20 个周期的训练,学习率为 0.001,没有学习率调度程序
  5. 使用交叉熵损失函数将像素分类为宠物、背景或宠物边界

在 20 个训练周期后,该模型实现了 86.96% 的验证精度。这比使用常规卷积的模型在相同的训练周期内实现的 88.28% 精度要低。我们已经通过实验确定,训练更多周期可以提高两个模型的准确性,因此 20 个周期肯定不是训练周期的终点。我们在本文中停止 20 个周期,以演示目的。

我们绘制了一个 gif,展示了该模型如何学习预测验证集中 21 张图像的分割掩模。

图6:一个 GIF,展示了使用深度可分离卷积的 SegNet 模型如何学习预测验证集中 21 张图像的分割掩模。来源:作者。

现在我们已经看到了模型在训练周期中的进展,让我们比较使用常规卷积和 DSC 的模型的训练周期。

准确性比较

我们发现查看使用常规卷积和 DSC 的模型的训练周期很有用。我们注意到的主要区别在于早期(周期)的训练阶段,之后两个模型大致都进入了相同的预测流程。事实上,在为两个模型训练 100 个周期后,我们注意到使用 DSC 的模型的准确度仅比使用常规卷积的模型低 1%。这符合我们仅训练 20 个周期的观察结果。

图7:使用常规卷积与 DSC 的 SegNet 模型预测分割掩模的进展。来源:作者。

您可能已经注意到,在仅进行 6 个训练周期后,两个模型都能够大致正确地预测出某些有用的内容。然后,大部分的模型训练工作就是确保预测掩模的边界尽可能紧密地贴近图像中的实际宠物边界。这意味着,虽然在后期训练周期中可以预期准确度的绝对增加较小,但其对预测质量的影响要大得多。我们注意到,在更高的绝对准确度值(从 89% 到 90%)上获得一位数字的准确度提高,会对预测产生显著的定量改进。

与UNet模型的比较

我们进行了一个实验,改变了很多超参数,重点是提高整体准确性,以了解此设置与最优设置相差多远。以下是该实验的配置。

  1. 图像大小:128 x 128 — 与目前的实验相同
  2. 训练周期:100 — 当前的实验训练了20个周期
  3. 增强:更多的增强,如图像旋转,通道丢失,随机块移除。我们使用Albumentations而不是torchvision transforms。 Albumentations会自动为我们转换分割掩模
  4. LR Scheduler:使用StepLR调度程序,每25个训练周期衰减0.8x
  5. 损失函数:我们尝试了4种不同的损失函数:交叉熵,聚焦,Dice,加权交叉熵。 Dice表现最差,而其他函数基本上相当。事实上,在100个周期之后,其余函数之间的最佳准确性差异在小数点后第4位(假设准确性是0.0到1.0之间的数字)
  6. 卷积类型:常规
  7. 模型类型:UNet — 当前的实验使用了SegNet模型

我们实现了上述设置的最佳验证准确性为91.3%。我们注意到图像大小显着影响最佳验证准确性。例如,当我们将图像大小更改为256 x 256时,最佳验证准确性提高到93.0%。但是,训练时间更长,使用更多内存,这意味着我们必须减少批量大小。

Figure 8: Result of training a UNet model for 100 train epochs with the hyperparameters mentioned above. Source: Author(s).

您可以看到与我们迄今为止看到的预测相比,这些预测要平滑得多,更清晰。

结论

在本系列的第3部分中,我们学习了深度可分离卷积(DSC)作为一种技术,用于减小模型大小和培训/推理成本,而不会显着损失验证准确性。我们了解了在特定设置下常规和DSC之间的预期大小/成本权衡。

我们展示了如何在PyTorch中适应SegNet模型以使用DSC。这种技术可以应用于任何深度CNN。实际上,我们可以有选择性地用DSC替换一些卷积层,即我们不需要一定替换所有卷积层。选择要替换的层将取决于您希望在模型大小/运行时成本和预测准确性之间取得的平衡。此决定将取决于您的具体用例和部署设置。

虽然本文对模型进行了20个周期的训练,但我们解释了这对于生产工作负载来说是不足的,并提供了一瞥,如果将模型训练更多个周期,将会发生什么。此外,我们提供了一些可以在模型训练期间调整的超参数的介绍。虽然这个列表并不全面,但它应该让您欣赏到为生产工作负载训练图像分割模型所需的复杂性和决策。

在本系列的下一部分中,我们将研究Vision Transformer以及如何使用此模型体系结构执行宠物分割任务的图像分割。

参考资料和进一步阅读

  • 高效深度学习书第04章——高效架构
  • 可分离卷积基础介绍