【论文精读】YOLO-World:实时开放词汇目标检测
论文地址: YOLO-World: Real-Time Open-Vocabulary Object Detection
源代码:YOLO-World
摘要
YOLO系列检测器因其高效性和实用性而被广泛认可。然而,它们依赖于预定义和训练过的物体类别,这限制了其在开放场景中的适用性。为了解决这一限制,我们提出了YOLO-World,这是一种创新的方法,通过视觉-语言建模和在大规模数据集上进行预训练,赋予YOLO开放词汇检测能力。具体来说,我们提出了一个新的可重参数化的视觉-语言路径聚合网络(RepVL-PAN)和区域-文本对比损失函数,以促进视觉和语言信息之间的交互。我们的方法在零样本检测中表现出色,能够高效地检测广泛的目标。在具有挑战性的LVIS数据集上,YOLO-World实现了35.4 AP的精度和52.0 FPS的速度,优于许多现有的先进方法。此外,经过微调的YOLO-World在多个下游任务中表现出色,包括目标检测和开放词汇实例分割。
1. Introduction
目标检测一直是计算机视觉中一个长期且基础性的重要挑战,它在图像理解、机器人技术和自动驾驶等领域有着广泛的应用。随着深度神经网络的发展,众多研究工作[16, 27, 43, 45]在目标检测领域取得了显著的突破。尽管这些方法取得了成功,但它们仍然存在局限性,因为它们只能处理具有固定词汇表(close-set)的目标检测任务,例如COCO数据集[26]中的80个类别。一旦目标类别被定义并标注后,训练好的检测器只能检测这些特定的类别,从而限制了其在开放场景中的能力和适用性。
最近的研究工作[8, 13, 48, 53, 58]探索了流行的视觉-语言模型[19, 39],以解决开放词汇检测[58]问题,通过从语言编码器(例如BERT[5])中提取词汇知识。然而,这些基于蒸馏的方法由于训练数据的稀缺性和词汇多样性的限制(例如OV-COCO[58]仅包含48个基础类别),其性能受到很大限制。一些方法[24, 30, 56, 57, 59]将目标检测训练重新定义为区域级视觉-语言预训练,并在大规模数据集上训练开放词汇目标检测器。然而,这些方法在真实世界场景中的检测仍然面临两个挑战:(1)计算负担重;(2)难以部署到边缘设备。以往的研究[24, 30, 56, 57, 59]已经展示了预训练大型检测器的潜力,而预训练小型检测器以赋予其开放识别能力仍然是一个未被探索的方向。
在本文中,我们提出了YOLO-World,旨在实现高效率的开放词汇目标检测,并探索大规模预训练方案,以将传统的YOLO检测器推向一个新的开放词汇世界。与以往方法相比,YOLO-World在推理速度和下游应用部署方面具有显著优势。具体而言,YOLO-World遵循标准的YOLO架构[20],并利用预训练的CLIP[39]文本编码器来编码输入文本。我们进一步提出了可重参数化的视觉-语言路径聚合网络(RepVL-PAN),以连接文本特征和图像特征,实现更好的视觉-语义表示。在推理阶段,可以移除文本编码器,并将文本嵌入重新参数化为RepVL-PAN的权重,以实现高效的部署。我们还通过在大规模数据集上进行区域-文本对比学习来研究YOLO检测器的开放词汇预训练方案,将检测数据、定位数据和图像-文本数据统一为区域-文本对。预训练的YOLO-World凭借丰富的区域-文本对,展现出强大的大规模词汇检测能力,更多的训练数据能够进一步提升其开放词汇能力。
图2. 与检测范式的比较。(a) 传统目标检测器:这些目标检测器只能检测训练数据集中预定义的固定词汇表内的目标,例如COCO数据集[26]中的80个类别。这种固定的词汇表限制了其在开放场景中的扩展能力。 (b) 以往的开放词汇检测器:以往的方法倾向于开发大型且重型的检测器用于开放词汇检测,这些检测器通常具有较强的能力。此外,这些检测器同时对图像和文本进行编码以进行预测,这在实际应用中非常耗时。 (c) YOLO-World:我们展示了轻量级检测器(例如YOLO检测器[20, 42])的强大开放词汇性能,这对于真实世界的应用具有重要意义。与使用在线词汇不同,我们提出了一个“提示-检测”范式以实现高效的推理。在这种范式中,用户根据需求生成一系列提示,并将这些提示编码为离线词汇表。然后,这些词汇表可以被重新参数化为模型权重,用于部署和进一步加速。
此外,我们探索了一种“提示-检测”范式,以进一步提高开放词汇目标检测在真实世界场景中的效率。如图2所示,传统的目标检测器[16, 20, 23, 41–43, 52]专注于固定词汇(close-set)检测,其类别是预先定义并训练好的。而以往的开放词汇检测器[24, 30, 56, 59]通过在线文本编码器对用户提示进行编码,以实现在线词汇检测。值得注意的是,这些方法倾向于使用大型检测器和重型骨干网络(例如Swin-L[32])来增强开放词汇能力。相比之下,我们的“提示-检测”范式(图2(c))首先对用户提示进行编码,以构建一个离线词汇表,该词汇表可以根据不同需求进行调整。然后,高效的检测器可以在推理时动态地检测离线词汇,而无需重新编码提示。对于实际应用,一旦我们训练好了检测器(即YOLO-World),我们可以预先对提示或类别进行编码,构建一个离线词汇表,并将其无缝集成到检测器中。
我们的主要贡献可以总结为以下三点:
- 我们提出了YOLO-World,这是一个具有高效率的先进开放词汇目标检测器,适用于真实世界的应用。
- 我们提出了一个可重参数化的视觉-语言路径聚合网络(RepVL-PAN),用于连接视觉和语言特征,并为YOLO-World设计了一个开放词汇的区域-文本对比预训练方案。
- 在大规模数据集上预训练的YOLO-World展现了强大的零样本性能,并在LVIS数据集上实现了35.4 AP的精度和52.0 FPS的速度。预训练的YOLO-World可以轻松适应下游任务,例如开放词汇实例分割和指代目标检测。此外,YOLO-World的预训练权重和代码将开源,以促进更多实际应用。
2. Related Works
2.1 传统目标检测
现有的目标检测研究主要集中在固定词汇表(close-set)检测上,其中目标检测器在具有预定义类别的数据集(例如COCO数据集[26]和Objects365数据集[46])上进行训练,然后仅检测固定类别中的目标。在过去几十年中,传统目标检测方法可以大致分为三类,即基于区域的方法、基于像素的方法和基于查询的方法。基于区域的方法[11, 12, 16, 27, 44],例如Faster R-CNN[44],采用两阶段框架进行提议生成[44]以及基于感兴趣区域(RoI)的分类和回归。基于像素的方法[28, 31, 42, 49, 61]倾向于采用单阶段检测器,这些检测器在预定义的锚点或像素上执行分类和回归。DETR[1]首次探索了基于Transformer[50]的目标检测,并启发了大量的基于查询的方法[64]。在推理速度方面,Redmon等人提出了YOLO系列[40–42],它们利用简单的卷积架构实现实时目标检测。一些工作[10, 23, 33, 52, 55]提出了YOLO的各种架构或设计,包括路径聚合网络[29]、跨阶段部分网络[51]和重参数化[6],这些改进进一步提升了速度和精度。与以往的YOLO系列相比,本文提出的YOLO-World旨在检测超出固定词汇表的目标,具备更强的泛化能力。
2.2 开放词汇目标检测
开放词汇目标检测(OVD)[58]作为现代目标检测的一个新趋势,旨在检测超出预定义类别的目标。早期工作[13]遵循标准的开放词汇设置[58],通过在基础类别上训练检测器并在评估时使用新的(未知的)类别。然而,尽管这种开放词汇设置可以评估检测器检测和识别新目标的能力,但由于训练数据集和词汇的有限性,它在开放场景中的泛化能力仍然受到限制。受视觉-语言预训练[19, 39]的启发,近期工作[8, 22, 53, 62, 63]将开放词汇目标检测重新定义为图像-文本匹配问题,并利用大规模图像-文本数据来扩展训练词汇。OWLViTs[35, 36]在检测和定位数据集上微调简单的视觉Transformer[7],构建了具有潜力的开放词汇检测器。GLIP[24]提出了一个基于短语定位的预训练框架,用于零样本设置下的开放词汇检测。Grounding DINO[30]将定位预训练[24]整合到检测Transformer[60]中,并引入跨模态融合。一些方法[25, 56, 57, 59]通过区域-文本匹配统一了检测数据集和图像-文本数据集,并在大规模图像-文本对上预训练检测器,取得了有希望的性能和泛化能力。然而,这些方法通常使用重型检测器(如ATSS[61]或DINO[60],并以Swin-L[32]作为骨干网络),导致计算需求高且部署困难。相比之下,我们提出了YOLO-World,旨在实现高效的开放词汇目标检测,具备实时推理能力,并且便于下游应用部署。与ZSD-YOLO[54]不同,后者也通过语言模型对齐探索YOLO的开放词汇检测,YOLO-World引入了一个新的YOLO框架和有效的预训练策略,增强了开放词汇性能和泛化能力。
3. Method
3.1 预训练公式化:区域-文本对
传统的目标检测方法(包括YOLO系列[20])是使用实例标注进行训练的,这些注释由边界框
和类别标签
组成。在本文中,我们将实例注释重新定义为区域-文本对
,其中文本可以是类别名称、名词短语或目标描述。YOLO-World采用图像和文本作为输入,并输出预测的边界框
和目标嵌入向量
,用于与输入文本中的类别或名词进行匹配。
图3. YOLO-World的整体架构。 与传统YOLO检测器相比,YOLO-World作为一种开放词汇检测器,采用文本作为输入。文本编码器首先将输入文本编码为文本嵌入。然后,图像编码器将输入图像编码为多尺度图像特征,提出的RepVL-PAN利用多级跨模态融合来处理图像和文本特征。最终,YOLO-World预测回归的边界框和对象嵌入,用于匹配输入文本中出现的类别或名词。
3.2 模型架构
YOLO-World的整体架构如图3所示,包括一个YOLO检测器、一个文本编码器和一个可重参数化的视觉-语言路径聚合网络(RepVL-PAN)。给定输入文本,文本编码器将文本编码为文本嵌入向量。YOLO检测器中的图像编码器从输入图像中提取多尺度特征。然后,我们利用RepVL-PAN通过图像特征和文本嵌入之间的跨模态融合来增强文本和图像的表示。
YOLO Detector. YOLO-World主要基于YOLOv8[20]开发,包含一个Darknet骨干网络[20, 43]作为图像编码器,一个路径聚合网络(PAN)用于多尺度特征金字塔,以及一个用于边界框回归和对象嵌入的头部。
Text Encoder. 给定文本输入,我们采用预训练的CLIP[39] Transformer文本编码器来提取文本嵌入向量。当输入文本是描述性语句或指代表达时,我们使用n-gram算法提取名词短语,并将其输入到文本编码器中。
Text Contrastive Head. 我们采用解耦head(包含两个3×3卷积)来回归边界框和对象嵌入向量
。我们设计了一个文text contrastive head,通过以下公式计算对象-文本相似度:
其中,L2-Norm(·)表示L2归一化,是第j个文本嵌入向量。我们还添加了仿射变换,其中α是可学习的缩放因子,β是平移因子。这些归一化和仿射变换对于稳定区域-文本训练至关重要。
Training with Online Vocabulary. 在训练阶段,我们为每个马赛克样本(包含4张图像)构建一个在线词汇表T,其中包含所有正样本名词和从数据集中随机采样的负样本名词。每个样本的词汇表最多包含M个名词,M默认设置为80。
Inference with Offline Vocabulary. 在推理阶段,我们提出了一种 prompt-then-detect 策略,使用离线词汇表以提高效率。用户可以根据需要定义一系列自定义提示,我们将这些提示编码为离线词汇嵌入向量。离线词汇表避免了每次输入的计算开销,并可以根据需要灵活调整词汇表。
图4. RepVL-PAN的示意图。本文提出的RepVL-PAN采用了文本引导的CSP层(Text-guided CSPLayer,简称T-CSPLayer),用于将语言信息注入图像特征中,同时还引入了图像池化注意力(Image Pooling Attention,简称I-Pooling Attention),用于增强图像感知的文本嵌入。
3.3. Re-parameterizable Vision-Language PAN
图4展示了RepVL-PAN的结构,它遵循自顶向下和自底向上的路径,通过多尺度图像特征构建特征金字塔
。我们提出了文本引导的CSP层(T-CSPLayer)和图像池化注意力(I-Pooling Attention),以增强图像特征和文本特征之间的交互,从而提高开放词汇检测的视觉-语义表示能力。在推理阶段,离线词汇嵌入向量可以被重参数化为卷积或线性层的权重,以便于部署。
Text-guided CSPLayer. 如图4所示,跨阶段部分层(CSPLayer)在自上而下或自下而上的融合后使用。我们通过将文本引导信息融入到多尺度图像特征中,扩展了[20]中的CSPLayer(也称为C2f),形成了文本引导的CSPLayer。具体来说,给定文本嵌入W和图像特征,我们在最后的深色瓶颈块之后采用最大sigmoid注意力机制,将文本特征与图像特征聚合,通过以下方式进行处理:
更新后的是与跨阶段特征拼接作为输出,δ表示Sigmoid函数。
Image-Pooling Attention.为了增强带有图像信息的文本嵌入,我们提出了图像池化注意力(Image-Pooling Attention),通过聚合图像特征来更新文本嵌入。不同于直接在图像特征上使用交叉注意力,我们对多尺度特征进行最大池化,以获得3×3的区域,从而生成27个补丁标记X̃ ∈ R27×D。然后,文本嵌入通过以下公式进行更新:
3.4. Pre-training Scheme
在本节中,我们介绍了在大规模检测、定位和图像-文本数据集上预训练YOLO-World的训练方案。
Learning from Region-Text Contrastive Loss. 给定马赛克样本I和文本T,YOLO-World输出K个对象预测,以及标注
。我们采用任务对齐标签分配[9]将预测与真实标注匹配,并为每个正预测分配一个文本索引作为分类标签。基于此词汇表,我们通过对象-文本相似度和对象-文本分配之间的交叉熵构建区域-文本对比损失
。此外,我们还采用IoU损失和分布式焦点损失进行边界框回归。总训练损失定义为:
其中,是一个指示因子,当输入图像I来自检测或定位数据时,
设置为1;当来自图像-文本数据时,
设置为0。由于图像-文本数据集中的边界框可能存在噪声,我们仅对具有准确边界框的样本计算回归损失。
Pseudo Labeling with Image-Text Data. 与直接使用图像-文本对进行预训练不同,我们提出了一种自动标签方法来生成区域-文本对。具体来说,标签方法包含三个步骤:(1)提取名词短语:我们首先利用n-gram算法从文本中提取名词短语;(2)伪标签:我们采用预训练的开放词汇检测器(例如,GLIP [24])为每个图像生成给定名词短语的伪框,从而提供粗略的区域-文本对;(3)筛选:我们使用预训练的CLIP [39]来评估图像-文本对和区域-文本对的相关性,并筛选低相关性的伪标注和图像。我们进一步通过整合非极大值抑制(NMS)等方法筛选冗余的边界框。我们建议读者参考附录中的详细方法。通过上述方法,我们从CC3M [47]中采样并标注了246k张图像,获得了821k个伪标注。
待更新