在美国等发达国家,道路每年变化高达 15%,Mapillary 通过将来自任何摄像机的图像组合成一个世界的 3D 可视化,满足了对地图更新日益增长的需求。Mapillary 的独立协作方法使任何人都可以收集、共享和使用街道级图像,以改进地图、发展城市和推动汽车行业。
如今,世界各地的人们和组织已贡献了超过 6 亿张图像,以支持 Mapillary 通过图像帮助人们理解世界各地的任务,并提供这些数据,其客户和合作伙伴包括世界银行、HERE 和丰田研究院。
Mapillary 的计算机视觉技术以前所未有的方式为地图带来了智能,增加了我们对世界的整体理解。Mapillary 大规模地对其所有图像进行最先进的语义图像分析和基于图像的 3D 建模。在这篇文章中,我们将讨论 Mapillary Research 的两项最新工作及其在 PyTorch 中的实现——无缝场景分割 [1] 和原地激活 BatchNorm [2]——分别生成全景分割结果并在训练期间节省高达 50% 的 GPU 内存。
无缝场景分割
Github 项目页面:https://github.com/mapillary/seamseg/

无缝场景分割的目标是从图像中预测“全景”分割 [3],即完整的标签,其中每个像素都分配有类别 ID,并在可能的情况下分配有实例 ID。像许多处理实例检测和分割的现代 CNN 一样,我们采用 Mask R-CNN 框架 [4],使用 ResNet50 + FPN [5] 作为主干。该架构分两个阶段工作:首先,“Proposal Head”在图像上选择一组可能包含对象的候选边界框(即提案);然后,“Mask Head”专注于每个提案,预测其类别和分割掩码。这个过程的输出是“稀疏”实例分割,只覆盖图像中包含可计数对象(例如汽车和行人)的部分。
为了完成我们的全景方法,即无缝场景分割,我们为 Mask R-CNN 添加了第三个阶段。源于相同的主干,“Semantic Head”对整个图像进行密集语义分割,同时考虑不可计数或无定形类别(例如道路和天空)。Mask 和 Semantic 头部的输出最终使用简单的非最大抑制算法融合以生成最终的全景预测。有关实际网络架构、使用的损失和底层数学的所有详细信息,请访问我们 CVPR 2019 论文 [1] 的项目网站。
虽然 Mask R-CNN 的几个版本是公开可用的,包括用 Caffe2 编写的官方实现,但在 Mapillary,我们决定从头开始使用 PyTorch 构建无缝场景分割,以便完全控制和理解整个管道。在此过程中,我们遇到了一些主要的绊脚石,并且不得不提出一些创造性的变通方法,我们将在下面进行描述。
处理可变大小张量
全景分割网络与传统 CNN 区别开来的一点是可变大小数据的普遍存在。事实上,我们处理的许多数量都无法轻易用固定大小的张量表示:每张图像包含不同数量的对象,Proposal head 可以为每张图像生成不同数量的提案,并且图像本身可以具有不同的大小。虽然这本身不是问题——可以一次处理一张图像——但我们仍然希望尽可能多地利用批处理级并行性。此外,在多 GPU 分布式训练时,DistributedDataParallel
期望其输入是批处理的、统一大小的张量。

我们解决这些问题的方法是将每批可变大小的张量封装在 PackedSequence
中。PackedSequence
不过是一个张量的美化列表类,将其内容标记为“相关”,确保它们都共享相同的类型,并提供有用的方法,例如将所有张量移动到特定设备等。当执行轻量级操作时,如果批处理级并行性不会快很多,我们只需在 for 循环中迭代 PackedSequence
的内容。当性能至关重要时,例如在网络的主体中,我们只需连接 PackedSequence 的内容,根据需要添加零填充(例如在具有可变长度输入的 RNN 中),并跟踪每个张量的原始维度。
PackedSequence
也有助于我们处理上面强调的第二个问题。我们稍微修改了 DistributedDataParallel
,使其能够识别 PackedSequence
输入,将其分成大小相等的块,并将它们的内容分发到 GPU 上。
与分布式数据并行器的非对称计算图
我们的网络另一个可能更微妙的特点是,它可以在 GPU 之间生成不对称的计算图。事实上,构成网络的一些模块是“可选的”,这意味着它们并非总是为所有图像计算。例如,当 Proposal head 不输出任何提案时,Mask head 就完全不会被遍历。如果我们使用 DistributedDataParallel
在多个 GPU 上进行训练,这会导致其中一个副本不计算 Mask head 参数的梯度。
在 PyTorch 1.1 之前,这会导致崩溃,因此我们不得不开发一种变通方法。我们简单但有效的解决方案是在不需要实际前向传播时计算一个“假前向传播”,即类似这样:
def fake_forward():
fake_input = get_correctly_shaped_fake_input()
fake_output = mask_head(fake_input)
fake_loss = fake_output.sum() * 0
return fake_loss
在这里,我们生成一批虚假数据,将其通过 Mask head,并返回一个始终将零反向传播到所有参数的损失。
从 PyTorch 1.1 开始,不再需要此变通方法:通过在构造函数中设置 find_unused_parameters=True
,DistributedDataParallel
会被告知识别未由所有副本计算梯度的参数并正确处理它们。这大大简化了我们的代码库!
原地激活 BatchNorm
Github 项目页面:https://github.com/mapillary/inplace_abn/
大多数研究人员可能会同意,无论他们的研究实验室拥有少量还是数千个 GPU,在可用的 GPU 资源方面总会存在限制。在 Mapillary 还在使用相对较少且主要是 12GB Titan X 风格的专业消费级 GPU 的时候,我们正在寻找一种解决方案,可以虚拟地增强训练期间可用的内存,以便我们能够在语义分割等密集标签任务上获得并推出最先进的结果。原地激活 BatchNorm 使我们能够使用多达 50% 的内存(计算开销很小),因此它已深度集成到我们目前的所有项目中(包括上面描述的无缝场景分割)。

在正向传播中处理 BN-激活-卷积序列时,大多数深度学习框架(包括 PyTorch)需要存储两个大的缓冲区,即 BN 的输入 x 和 Conv 的输入 z。这是必要的,因为 BN 和 Conv 的反向传播的标准实现依赖于它们的输入来计算梯度。使用 InPlace-ABN 替换 BN-激活序列,我们可以安全地丢弃 x,从而在训练时节省高达 50% 的 GPU 内存。为了实现这一点,我们根据 BN 的输出 y 重写了 BN 的反向传播,而 y 又通过反转激活函数从 z 重构。
InPlace-ABN 唯一的限制是它需要使用可逆激活函数,例如 Leaky ReLU 或 ELU。除此之外,它可以在任何网络中直接替换 BN+激活模块。我们的原生 CUDA 实现与 PyTorch 的标准 BN 相比提供了最小的计算开销,并且可供任何人从这里使用:https://github.com/mapillary/inplace_abn/。
具有非对称图和不平衡批次的同步 BN
在多个 GPU 和/或多个节点上训练同步 SGD 网络时,通常的做法是在每个设备上单独计算 BatchNorm 统计信息。然而,根据我们使用语义和全景分割网络的工作经验,我们发现累积所有工作器的均值和方差可以显著提高准确性。在处理小批量时尤其如此,例如在无缝场景分割中,我们每个 GPU 只训练一张超高分辨率图像。
InPlace-ABN 支持在多个 GPU 和多个节点上的同步操作,并且从 1.1 版开始,这也可以在标准 PyTorch 库中使用 SyncBatchNorm 实现。然而,与 SyncBatchNorm 相比,我们支持一些额外的功能,这对于无缝场景分割尤为重要:不平衡批次和非对称图。
如前所述,Mask R-CNN 类网络自然会产生可变大小的张量。因此,在 InPlace-ABN 中,我们使用此处描述的并行算法的变体来计算同步统计信息,该算法正确考虑了每个 GPU 可以容纳不同数量样本的事实。PyTorch 的 SyncBatchNorm 目前正在修订以支持此功能,改进后的功能将在未来版本中提供。
非对称图(如上所述)是创建同步 BatchNorm 实现时必须处理的另一个复杂因素。幸运的是,PyTorch 的分布式组功能允许我们将分布式通信限制在工作器子集,轻松排除当前不活动的工作器。唯一缺少的部分是,为了创建一个分布式组,每个进程都需要知道将参与该组的所有进程的 ID,即使不属于该组的进程也需要调用 new_group()
函数。在 InPlace-ABN 中,我们用这样的函数来处理它:
import torch
import torch.distributed as distributed
def active_group(active):
"""Initialize a distributed group where each process can independently decide whether to participate or not"""
world_size = distributed.get_world_size()
rank = distributed.get_rank()
# Gather active status from all workers
active = torch.tensor(rank if active else -1, dtype=torch.long, device=torch.cuda.current_device())
active_workers = torch.empty(world_size, dtype=torch.long, device=torch.cuda.current_device())
distributed.all_gather(list(active_workers.unbind(0)), active)
# Create group
active_workers = [int(i) for i in active_workers.tolist() if i != -1]
group = distributed.new_group(active_workers)
return group
首先,每个进程(包括不活动的进程)通过 all_gather
调用将其状态传达给所有其他进程,然后它使用共享信息创建分布式组。在实际实现中,我们还为组添加了缓存机制,因为 new_group()
通常太昂贵,无法在每个批次都调用。
参考文献
[1] 无缝场景分割;Lorenzo Porzi, Samuel Rota Bulò, Aleksander Colovic, Peter Kontschieder;计算机视觉与模式识别(CVPR),2019
[2] 用于 DNN 内存优化训练的原地激活 BatchNorm;Samuel Rota Bulò, Lorenzo Porzi, Peter Kontschieder;计算机视觉与模式识别(CVPR),2018
[3] 全景分割;Alexander Kirillov, Kaiming He, Ross Girshick, Carsten Rother, Piotr Dollar;计算机视觉与模式识别(CVPR),2019
[4] Mask R-CNN;Kaiming He, Georgia Gkioxari, Piotr Dollar, Ross Girshick;国际计算机视觉会议(ICCV),2017
[5] 用于目标检测的特征金字塔网络;Tsung-Yi Lin, Piotr Dollar, Ross Girshick, Kaiming He, Bharath Hariharan, Serge Belongie;计算机视觉与模式识别(CVPR),2017