多任务架构：全面指南

轻量级模型用于实时多任务推理

介绍

您是否曾经想过如何训练一个深度神经网络来完成多项任务？这样的模型被称为多任务架构，与为每个任务使用单独模型的传统方法相比，具有一定的优势。多任务架构是多任务学习的子集，多任务学习是训练模型或一组模型同时执行多个任务的一种通用方法。

在本文中，我们将学习如何训练一个单一模型同时执行分类和回归任务。本文的代码可以在GitHub上找到。以下是概述：

• 动机 — 为什么要这样做？
• 方法 — 我们将如何做到这一点？
• 模型架构
• 训练方法
• 推理 — 检查性能并从一个有趣的失败中学习
• 结论

动机

为什么我们要使用轻量级模型？这难道不会降低性能吗？如果我们不部署到边缘设备，我们不应该尽可能使用大型模型吗？

边缘应用需要轻量级模型以实现低功耗的实时推理。其他应用也可以从中受益，但是如何受益呢？轻量级模型被低计算要求所忽视。一般来说，这可以降低服务器使用率，从而降低功耗。这将总体上降低成本并减少碳排放，后者在人工智能的未来可能成为一个重大问题。

通过减少功耗，轻量级模型可以帮助降低成本和减少碳排放

话虽如此，多任务架构只是工具箱中的一个工具，在决定使用哪些工具之前，应考虑所有项目要求。现在让我们深入了解如何训练其中之一的示例！

方法

为了构建我们的多任务架构，我们将大致涵盖这篇论文中的方法，该论文训练了一个用于同时分割和深度估计的单一模型。其基本目标是以快速高效的方式执行这些任务，而性能的损失是可以接受的。在多任务学习中，我们通常将相似的任务分组在一起。在训练过程中，我们还可以添加一个辅助任务，可能有助于模型的学习，但我们可能决定在推理过程中不使用它[1, 2]。为简单起见，我们在训练过程中不使用任何辅助任务。

深度和分割都是密集预测任务，并且具有相似性。例如，单个对象的深度可能在对象的所有区域都保持一致，形成一个很窄的分布。主要思想是每个单独的对象应具有自己的深度值，并且我们应该能够通过查看深度图像仅仅识别出单个对象。同样地，我们应该能够通过查看分割图像识别出相同的单个对象。虽然可能会有一些异常值，但我们将假设这种关系成立。

数据集

我们将使用City Scapes数据集提供(左摄像头)输入图像分割掩码和深度图。对于分割图像，我们选择使用标准的训练标签，其中包含19个类别+ 1个未标记类别。

深度图准备 — 默认视差

使用SteroSGBM创建的视差图可以从CityScapes网站上轻松获取。视差描述了从每个立体相机的视角观察到的对象之间的像素差异，它与深度成反比，可以使用以下公式计算：

然而，默认的视差图存在很多噪点，许多空洞对应无限深度以及始终显示自车的部分。清理这些视差图的常见方法包括：

1. 剪裁底部的20%，以及左侧和顶部边缘的部分
2. 恢复到原始比例
3. 应用平滑滤波器
4. 进行修复

一旦我们清理了视差图，就可以计算深度，结果如下：

这种方法的详细细节超出了本篇文章的范围，但是如果您感兴趣，可以在YouTube上观看视频解释。

剪裁和调整大小的步骤意味着视差（以及深度）图与输入图像不完全对齐。尽管我们可以对输入图像进行相同的剪裁和调整大小来纠正这一点，但我们选择探索一种新方法。

深度图准备 – CreStereo视差

我们探索使用CreStereo从左右图像生成高质量的视差图。CreStereo是一个先进的模型，能够预测出平滑的视差图。这种方法引入了一种被称为知识蒸馏的范式，其中CreStereo是一个教师网络，而我们的模型将是学生网络（至少用于深度估计）。这种方法的详细细节超出了本篇文章的范围，但是如果您感兴趣，可以在YouTube上观看此链接。

通常，CreStereo的深度图噪声很小，因此不需要剪裁和调整大小。然而，分割蒙版中存在的自车可能会导致泛化问题，因此在所有训练图像上移除了底部的20%。下面显示了一个训练样本：

现在我们有了数据，让我们看看架构。

模型架构

根据[1]，架构将包括MobileNet主干/编码器、轻量级RefineNet解码器和每个单独任务的头部。整体架构如下图所示（图3）。

对于编码器/主干，我们将使用预训练的MobileNetV3，并将1/4、1/8、1/16和1/32分辨率的跳跃连接传递给轻量级Refine Net。最后，输出传递给负责不同任务的每个头部。请注意，如果需要，我们甚至可以向此架构添加更多任务。

为了实现编码器，我们使用预训练的MobileNetV3编码器，将MobileNetV3编码器传递给自定义的PyTorch模块。它的forward函数的输出是用于输入到轻量级Refine Net的跳跃连接的ParameterDict。下面的代码片段展示了如何实现这一点。

class MobileNetV3Backbone(nn.Module):    def __init__(self, backbone):        super().__init__()        self.backbone = backbone        def forward(self, x):        """ Passes input theough MobileNetV3 backbone feature extraction layers            layers to add connections to                - 1:  1/4 res                - 3:  1/8 res                - 7, 8:  1/16 res                - 10, 11: 1/32 res           """        skips = nn.ParameterDict()        for i in range(len(self.backbone) - 1):            x = self.backbone[i](x)            # add skip connection outputs            if i in [1, 3, 7, 8, 10, 11]:                skips.update({f"l{i}_out" : x})        return skips

轻量级RefineNet解码器与[1]中实现的解码器非常相似，只是对其进行了一些修改，以使其与MobileNetV3兼容，而不是MobileNetV2。我们还注意到解码器部分由分割和深度两个部分组成。模型的完整代码可在GitHub上找到。我们可以按照以下方式组装模型：

from torchvision.models import mobilenet_v3_small    mobilenet = mobilenet_v3_small(weights='IMAGENET1K_V1')encoder = MobileNetV3Backbone(mobilenet.features)decoder = LightWeightRefineNet(num_seg_classes)model = MultiTaskNetwork(encoder, freeze_encoder=False).to(device)

训练方法

我们将训练分为三个阶段，第一个阶段为1/4分辨率，第二个阶段为1/2分辨率，最后一个阶段为完整分辨率。由于冻结编码器权重似乎不能产生良好的结果，所有权重都被更新。

变换

在每个阶段中，我们执行随机裁剪调整大小、颜色抖动、随机翻转和归一化。左侧输入图像使用标准图像均值和标准差进行归一化。

深度转换

一般来说，深度图中包含的大部分值较小，因为深度图中包含的大部分信息是相机附近的物体和表面。由于深度图的深度主要集中在较低的值周围（见下图4左侧），它需要进行转换才能被神经网络有效地学习。深度图被裁剪在0到250之间，这是因为大距离处的立体视差/深度数据通常不可靠，而在这种情况下，我们希望有一种方法来丢弃它。然后，我们取其自然对数并将其除以5，以将分布压缩在较小的数字范围内。有关更多详细信息，请参见此笔记本。

老实说，我不确定如何最好地转换深度数据。如果有更好的方法或者您有不同的做法，我会很乐意在评论中了解更多：）。

损失函数

我们使用简单的交叉熵损失进行分割，使用均方差损失进行深度估计。我们将它们加在一起，没有权重，并同时进行优化。

学习率

我们使用一种带有最大值为5e-4的单周期余弦退火学习率，在1/4分辨率下进行150个周期的训练。用于训练的笔记本位于此处。

我们在1/2分辨率下进行25个周期的微调，然后在完整分辨率下再进行25个周期的微调，学习率均为5e-6。请注意，每次在增加分辨率时微调时，我们需要减小批量大小。

推理

对于推理，我们对输入图像进行归一化，并通过模型进行正向传递。图6显示了来自验证集和测试集的训练结果。

总体来说，当图像中有较大的物体时，模型能够对其进行分割和估计深度。当存在细节更加复杂的物体，比如行人时，模型往往难以完全分割它们。模型能够在一定程度上估计它们的深度。

一个有趣的失败案例

图6的底部展示了一个有趣的失败案例，未能完全分割图像左侧的灯杆。分割只覆盖了灯杆的下半部分，而深度显示，灯杆的下半部分比上半部分更接近。深度失败可能是由于底部像素偏向于较近深度的偏见导致的；注意像素500周围的地平线，较近的像素和较远的像素之间存在明显的分界线。看起来这种偏见可能已经泄漏到了模型的分割任务中。在训练多任务模型时，应考虑到这种任务泄漏的情况。

在多任务学习中，一个任务的训练数据可能会影响另一个任务的性能

深度分布

让我们来看看预测的深度与真实值的分布情况。为了简单起见，我们将只使用94个真实/预测的完整分辨率深度图对。

看起来模型学习到了两个分布，一个在4附近有一个峰值，另一个在30附近有一个峰值。注意到剪切伪影似乎没有产生影响。整体分布包含了一个长尾，这是因为图像中只有一小部分包含远处的深度数据。

预测的深度分布比真实值更加平滑。真实值分布的粗糙性可能源于每个物体包含相似的深度值。可以尝试利用这些信息来应用某种正则化，以强制模型遵循这种范式，但这将留待以后再讨论。

额外内容：推理速度

由于这是一个旨在提高速度的轻量级模型，让我们看看它在GPU上的推理速度如何。下面的代码已经修改自这篇文章。在这个测试中，输入图像已被缩小到400×1024。

# 寻找执行卷积的最佳后端torch.backends.cudnn.benchmark = True# 缩小到一半大小rescaled_sample = Rescale(400, 1024)(sample)rescaled_left = rescaled_sample['left'].to(DEVICE)# 初始化记录器starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)repetitions = 300timings=np.zeros((repetitions,1))# GPU热身for _ in range(10):    _, _ = model(rescaled_left.unsqueeze(0))# 测量性能with torch.no_grad():    for rep in range(repetitions):        starter.record()        _, _ = model(rescaled_left.unsqueeze(0))        ender.record()        # 等待GPU同步torch.cuda.synchronize()        curr_time = starter.elapsed_time(ender)        timings[rep] = curr_timemean_syn = np.sum(timings) / repetitionsstd_syn = np.std(timings)print(mean_syn, std_syn)

推理测试显示，该模型可以以18.69+/-0.44毫秒的速度运行，约为55Hz。需要注意的是，这只是在配备NVIDIA RTX 3060 GPU的笔记本电脑上运行的Python原型，不同的硬件将改变推理速度。我们还应该注意，像Torch-TensorRt这样的SDK如果部署在NVIDIA GPU上，可以提供显著的加速。

结论

在本文中，我们学习了多任务学习如何节省成本和减少碳排放。我们学习了如何构建一个轻量级多任务架构，能够在CityScapes数据集上同时执行分类和回归。我们还利用了CreStereo和知识蒸馏来帮助我们的模型学习预测更好的深度图。

这个轻量级模型在速度和效率方面做出了权衡。即使在这种权衡下，训练好的模型在测试数据上能够预测出合理的深度和分割结果。此外，它还能够学习到与真实深度图类似的深度分布。

参考文献

[1] Nekrasov, Vladimir, et al. ‘使用非对称注释进行实时联合语义分割和深度估计’. CoRR, vol. abs/1809.04766, 2018, http://arxiv.org/abs/1809.04766

[2] Standley, Trevor, et al. ‘多任务学习中应该一起学习哪些任务？’ CoRR, vol. abs/1905.07553, 2019, http://arxiv.org/abs/1905.07553

[3] Cordts, M., Omran, M., Ramos, S., Rehfeld, T., Enzweiler, M., Benenson, R., Franke, U., Roth, S., & Schiele, B. (2016). 用于语义城市场景理解的城市景观数据集。2016 IEEE计算机视觉与模式识别（CVPR）会议。https://doi.org/10.1109/cvpr.2016.350