重新审视模型架构的掩码头部,用于新类别实例分割

发布人:Google Research 研究软件工程师 Vighnesh Birodkar 和研究员 Jonathan Huang

实例分割的任务是将图像中的像素分组为单个事物的实例,并用类标签(可计数的对象,如人、动物、汽车等,并为每个对象分配独特的标识符,如 car_1 和 car_2)来标识这些事物。实例分割作为一项核心的计算机视觉任务,对许多下游应用至关重要,如自动驾驶汽车、机器人、医学成像和照片编辑。

近年来,深度学习 (Deep learning) 在使用 Mask R-CNN 之类的架构解决实例分割问题方面取得了重大进展。然而,这些方法依赖于收集大型的标签实例分割数据集。但不同于收集边界框标签的方法,如 Extreme clicking,可以实现每个实例 7 秒的收集速度,收集实例分割标签(称为“掩码”)时,每个实例用时可能需要 高达 80 秒,该方式较高的成本,拉高了这项研究的门槛。另一个相关任务——泛型分割,甚至需要更多的标签数据。

部分监督的实例分割环境(即只用实例分割掩码给一小部分类加标签,其余大部分类只用边界框来加标签)这一方法有可能减少对人工创建的掩码标签的依赖,从而大大降低开发实例分割模型的门槛。不过,这种部分监督的方法也需要更强的模型泛化形式来处理训练时没有遇到过的新类别,例如,只用动物掩码进行训练,然后让模型针对建筑物或植物产生准确的实例分割。此外,还有简单的方法,例如训练一个与类无关的 Mask R-CNN,同时忽略任何没有掩码标签的实例的掩码损失 (Loss function),但这些方法效果并不好。例如,在典型的 “VOC/Non-VOC” 基准中,Mask R-CNN 针对 COCO 中 20 个类的子集(称为“已见类”)进行掩码训练,并在其余 60 个类(称为“未见类”)上进行测试,一个带有 Resnet-50 主干的典型 Mask R-CNN 在未见类上的 掩码 mAP(即平均精度,数值越高越好)只能达到约 18%,而在全监督时,在同一集合上的掩码 mAP 则高出很多,超过了 34%。

在即将发布于 ICCV 2021 的 “掩码头部架构对新类别分割的惊人影响 (The surprising impact of mask-head architecture on novel class segmentation)”一文中,我们确定了 Mask R-CNN 在新类别上表现不佳的主要原因,并提出了两个易于实施的修复方法(训练协议修复;掩码头部架构修复),这两种方法协同作用,可以缩小与全监督性能之间的差距。

我们证明了这种方法普遍适用于裁剪-分割模型,即 Mask R-CNN 或类似 Mask R-CNN 的架构:计算整个图像的特征表征,然后将每个实例的裁剪传递给第二阶段的掩码预测网络(也称为掩码头部网络)。对发现结果进行整合,我们提出了基于 Mask R-CNN 的模型,该模型的掩码 mAP 远高于目前最先进的模型,提升了 4.7%,且无需更复杂的辅助损失函数、离线训练的先验因素或先前研究中提出的权重转移函数。我们还开放了该模型两个版本的代码库,分别称为 Deep-MACDeep-MARC,并发布了一个 colab,从而以互动方式生成掩码,如下面的视频演示所示。

模型 DeepMAC 的演示版,即使是训练时未见过的类,该模型也可以学习预测准确的掩码,以及给定用户所指定的边框。亲自在 colab 中试试吧。图片来源:Chris Briggs、维基百科和 Europeana

部分监督环境中裁剪方法的影响

裁剪是裁剪-分割模型的一个重要步骤,通过裁剪特征图以及对应每个实例的边界框的实际掩码来训练 Mask R-CNN。将这些裁剪过的特征传递给另一个神经网络(称为掩码头部网络),该网络计算出最终的掩码预测,然后将其与掩码损失函数中的实际裁剪进行比较。裁剪有两种选择:(1) 直接对实例的实际边界框进行裁剪, (2) 对模型预测的边界框(称为建议)进行裁剪。在测试时始终通过建议来执行裁剪,因为要假设实际边界框不可用。

“对实际边界框的裁剪”对比“训练过程中对模型预测的建议裁剪”。标准 Mask R-CNN 实现使用上述两种类型的裁剪,但我们已经证明,只对实际边界框进行裁剪在新类别上表现出显著的性能优势

我们考虑了一个类似于 Mask R-CNN 的一般架构系列,与典型的 Mask R-CNN 训练环境相比,存在一个微小但关键的区别:我们在训练时使用实际边界框(而不是建议边界框)裁剪

典型的 Mask R-CNN 实现将两种类型的裁剪都传递给掩码头部。然而,在传统的观点中,这个选择是一个不重要的实施细节,因为它在全监督环境中不会对性能产生显著影响。相反,对于部分监督环境,我们发现裁剪方法起着重要的作用,虽然在训练过程中,只对实际边界框进行裁剪不会使全监督环境下的结果发生明显变化,但在部分监督环境中却有着惊人的显著积极影响,在未见类上的表现明显改善。

利用建议和实际边界框(默认设置)或只用实际边界框进行训练时,Mask R-CNN 在未见类上的性能。只用实际边界框训练掩码头部时,在未见类上的性能有明显的提升,mAP 超过 9%。我们报告了 ResNet-101-FPN 主干加持下的性能

解锁掩码头部的完全泛化潜力

更令人惊讶的是,上述方法引发了一个新现象:在训练过程中启用实际裁剪, Mask R-CNN 的掩码头部对模型的泛化能力(泛化至未见类)起着异常重要的作用。举个例子,我们在下图中比较了几个模型,对象为停车计时器、手机和披萨(训练期间未见过的类)。每个模型都已启用实际边界框裁剪,但使用的开箱即用掩码头部架构不同。

使用四种不同的掩码头部架构对未见类进行掩码预测(从左到右分别是:ResNet-4ResNet-12ResNet-20Hourglass-20,其中数字是指神经网络的层数)尽管从未见过“停车计时器”、“披萨”或“手机”类中的掩码,但最右的掩码头部架构可以正确分割这些类。我们展示的掩码头部架构在掩码预测方面的性能从左到右依次递增。此外,这种差异只有在未见类上进行评估时才比较明显,如果在已见类上进行评估,所有四个架构会表现出类似的性能

特别需要注意的是,在全监督环境中,掩码头部架构之间在这些方面的差异并不明显。顺便说一下,这可能解释了为什么先前的实例分割研究几乎只使用浅层(即低层数)掩码头部,因为增加复杂性无法带来任何优势。下面我们比较了三种不同的掩码头部架构在已见与未见类上的掩码 mAP。所有这三种模型在已见类的集合上展现了同样优越的性能,但应用于未见类时,深 沙漏型 掩码头部脱颖而出。我们发现,在所尝试的架构中,沙漏型掩码头部效果是最好的,并且在使用 50 层以上的沙漏型掩码头部时获得了最佳结果。

ResNet-4、Hourglass-10 和 Hourglass-52 掩码头部架构在已见类和未见类上的性能。尽管在已见类上的性能几乎没有变化,但在未见类上的性能却有很大差别

最后,我们证明这一发现具有普遍性,适用于各种主干(如 ResNet、SpineNet 和 Hourglass)和检测器架构,包括 基于锚无锚的检测器,甚至在根本没有检测器的情况下也适用。

总结

为了得出最完善的结果,我们整合了上述发现:我们在高分辨率图像 (1280x1280) 上利用 SpineNet 主干训练了一个启用实际边界框裁剪且带有深 Hourglass-52 掩码头部的 Mask R-CNN 模型。我们称此模型为 Deep-MARC (Deep Mask heads Above R-CNN)。在不使用任何离线训练或其他手动先验因素的情况下,Deep-MARC 超过了之前最先进的模型,掩码 mAP 提高了 4.5%(绝对值)以上。为证明这种方法的普遍性,我们还训练了基于 CenterNet(而非基于 Mask R-CNN)的模型(称为 Deep-MAC),该模型同样展现出强大的性能,也超越了之前最先进的水平。

Deep-MAC 和 Deep-MARC 与其他部分监督实例分割方法(如 MaskX R-CNN、ShapeMask 和 CPMask)的比较

结论

我们开发的实例分割模型能够泛化到不属于训练集的类。这其中要强调两个关键因素的作用,这两个因素可以应用于任何裁剪-分割模型(如 Mask R-CNN):(1) 训练过程中的实际边界框裁剪, (2) 强大的掩码头部架构。虽然这两个因素对训练期间有掩码的类影响不大,但在训练期间没有掩码的新类别上,采用这两个因素会带来明显的改善。此外,这两个因素足以在部分监督的 COCO 基准上实现最先进的性能。最后,我们的研究结果具有普遍性,也可能对相关的任务产生影响,如 全景分割姿势预测

致谢

感谢联合作者 Zhichao Lu、Siyang Li 和 Vivek Rathod。感谢 David Ross 和匿名的 ICCV 审核者,他们为改善此项研究提供了至关重要的评论意见。

原文:Revisiting Mask-Head Architectures for Novel Class Instance Segmentation
中文:TensorFlow 公众号