- Date:2024.02.22,Tencent AI Lab,华中科技大学
- Paper:https://arxiv.org/pdf/2401.17270.pdf
- Github:https://github.com/AILab-CVC/YOLO-World
论文解决的问题: 通过视觉语言建模和大规模数据集上的预训练来增强YOLO的开发词汇检测能力。
YOLO-world提出了一种prompt-then-detect范式:先提示,再检测。即将单词先转化成一系列离线的embedding,再将其重参数为模型的参数,参与到目标检测任务中来。相当于在部署的时候可以直接砍掉CLIP部分,直接把想要检测的类别的text embedding放到YOLO-World中进行推理,所以前向速度非常快。
文章目录
- 1 文章贡献
- 2 网络结构
- 2.1 RepVL-PAN(Reparameterizable Vision-Language PAN)
- 2.2 Text Contrastive Head
- 3 训练过程
- 3.1 损失函数
- 3.2 训练策略
- 4 训练数据的生成
- 5 消融实验上的一些结论
1 文章贡献
- 推出了一个可部署的开放词汇目标检测器YOLO-World,可用于部署。具有较强的零样本能力,在LVIS数据集上的map指标为35.4,fps为52(V100);
- 提出了一个可重参数化的网络的结构RepVL-PAN来连接视觉和文本特征;
2 网络结构
- text encoder使用的预训练的CLIP的text encoder,当文本是描述时,使用n-gram算法提取名词
- 文本的特征于图像的特征(multi-scale)在RepVL-PAN结构中进行融合
- RepVL-PAN输出Image-aware的text embedding,检测段输出bbox和text-aware的image embedding
- 最后再通过计算text embedding和image embedding之间距离来判断目标框中的类别
- YOLO backbone使用的是YOLOv8
2.1 RepVL-PAN(Reparameterizable Vision-Language PAN)
-
Text-guide CSPLayer: 采用max-sigmoid attention来实现文本和图像的融合其中 X l X_l Xl是来自不同层的图像特征,W是text embedding,计算特征图与text embedding的相似度矩阵,取最大值在sigmoid后作为权重加权原特征图
X l ′ = X l ⋅ δ ( max j ∈ { 1.. C } ( X l W j ⊤ ) ) ⊤ X_l^{\prime}=X_l \cdot \delta\left(\max _{j \in\{1 . . C\}}\left(X_l W_j^{\top}\right)\right)^{\top} Xl′=Xl⋅δ(j∈{1..C}max(XlWj⊤))⊤
不过在代码实现的时候,其实concat了每一个block的输出:
# yolo_world/models/layers/yolo_bricks.py:145行 def forward(self, x: Tensor, guide: Tensor) -> Tensor:"""Forward process."""x_main = self.main_conv(x)x_main = list(x_main.split((self.mid_channels, self.mid_channels), 1))x_main.extend(blocks(x_main[-1]) for blocks in self.blocks)x_main.append(self.attn_block(x_main[-1], guide))return self.final_conv(torch.cat(x_main, 1))
-
Image-Pooling Attention: 利用图像感知信息增强text embedding,文本嵌入的更新方式如下:
W ′ = W + MultiHead-Attention ( W , X ^ , X ^ ) W^{\prime}=W+\text{MultiHead-Attention}(W,\hat X, \hat X) W′=W+MultiHead-Attention(W,X^,X^)
X是来自于3个不同尺度的图像特征,对每个尺度进行最大池化max-pooling,得到3x3大小的特征图,3组就是27个patch-token。
2.2 Text Contrastive Head
YOLO-World的检测头还是复用了YOLOv8的解耦头,只不过将分类分支修改成了object embedding。object embedding与text embedding需要先进行L2范数归一化,再进行距离计算,其计算公式如下:
s k , j = α ∣ ∣ e k ∣ ∣ 2 ∗ ∣ ∣ w j ∣ ∣ 2 + β s_{k,j} = \alpha ||e_k||^2*||w_j||^2 + \beta sk,j=α∣∣ek∣∣2∗∣∣wj∣∣2+β
其中 s k , j s_{k,j} sk,j表示文本 w j w_j wj与目标嵌入 e k e_k ek之间的距离, α \alpha α和 β \beta β是可学习的比例因子和位移因子。
3 训练过程
3.1 损失函数
- 输出K个预测结果和其对应的gt:{box,text},使用TaskAlignAssigner进行匹配。匹配后通过计算object-text之间的交叉熵,也就是所谓的Region-text constrastive loss(没错,就是CLIP中Image-text constrastive loss,只不过这里不是使用全图的embedding,而是使用检测区域的embedding)。
- 使用IoU loss和DFL(Distribute Focal Loss)来计算box的回归损失,这就意味着yolo-world使用的也是解耦头,并将回归任务转化成了分类任务。最后的总loss为:
L ( I ) = L c o n + λ i ( L i o u + L d f l ) L(I) = L_{con} + \lambda_i(L_{iou} + L_{dfl}) L(I)=Lcon+λi(Liou+Ldfl)
其中 λ i \lambda_i λi是indicator factor(指标因子),当图像来自于detection或者grounding data时设置为1,来自于image-text时设置为0。
3.2 训练策略
训练使用在线词汇表,推理使用离线词汇表。
-
训练时:
- 输入图片:4张组成的mosaic图片
- 在线词汇表:T
- 每次训练的时候需要从词汇表中抽取图片中含有的N个名词,再随机抽选80-N个不存在于该图中的目标名称,因此每次送入网络中的名词数量默认为80个
-
测试时: 设置需要的词,并直接获取词向量参与到目标检测中,而不再用CLIP进行编码。
-
其他细节:
- 训练框架:MMYOLO
- Text encoder:预训练的CLIP
- GPU:32个V100
- Batch:512
- 数据增强:随机色彩、翻转、仿射、4张mosaic
4 训练数据的生成
文章中设计了一个数据生产流程,对CC3M中的246K图片生成了821K的伪标注。
其生成流程如下:
- 提取名词短语: n-gram
- 伪标签(pseudo label): 使用预训练的开放词汇检测器(GLIP),给每个图像给定的名词短语生成检测框,从而提供粗略的region-text pairs
- 过滤(filter): 使用预训练的CLIP来评估region-text pairs的相关性,过滤掉相关性低的注释和图像。给定图像I,图像描述T和粗略的region-text对 ( B i , t i , c i ) (B_i,t_i,c_i) (Bi,ti,ci)
3.1. 计算Image-text score得到S(img)
3.2. 裁出region区域,计算region-text的相似度S(region)
3.3. [可选] Re-labeling:裁剪后的图片与所有名词计算相似度,选择最高的名词作为该区域的text,用于纠正GLIP错标的文本
3.4. Rescoring:用region-text的相似度S(region)对置信度重新评分 c i = ( c i ∗ s i r e g i o n ) c_i = \sqrt{(c_i * s^{region}_i)} ci=(ci∗siregion)
3.5. 区域级过滤(Region-level filtering):根据文本分成不同的组,同类别使用NMS(0.5),过滤置信度低的(0.3)
3.6. 图像级过滤(Image-level filtering):计算保留的所有region的平均分数,得到图像级置信度分数 S = ( S ( i m g ) ∗ S ( r e g i o n ) ) S = \sqrt{(S(img) * S(region))} S=(S(img)∗S(region)),保留分数大于0.3的图像。 - 使用NMS过滤冗余的检测框。
- 训练集数据情况
- 测试集数据情况: 在LVIS进行zero-shot evaluation
- 包含1203个类别,远超训练集的类别
- 主要在LVIS minival上测试固定AP
- 最大预测数量设置为1000
总结:可以看出检测效果非常Amazing,而且更关键的是速度非常快! 其速度快的原因在于测试的时候,可以直接将text embeding先算出来,让后作为YOLO-World的参数来进行前向推理。直接砍掉了CLIP在推理阶段的耗时。这个思路非常nice!
5 消融实验上的一些结论
结论一:增加高质量数据效果提升显著。
验证Text-guided CSPLayers和Image-Pooling Attention的作用分别可以提升0.8和0.3个点,但是只加入GQA数据集可以直接提升7.2个点,还是证明了加入数据的有效性。
*注:GQA是斯坦福大学教授 Christopher Manning 及其学生 Drew Hudson 一同打造的全新图像场景图问答数据集,旨在推动场景理解与视觉问答研究领域的进步,详细数据链接:https://cs.stanford.edu/people/dorarad/gqa * 。
结论二:在YOLO-Wolrd里面,CLIP优于BERT。