【论文阅读】YOLO-World | 开集目标检测

• 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 文章贡献

• 推出了一个可部署的开放词汇目标检测器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$是来自不同层的图像特征，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​(Xl​Wj⊤​))⊤
  不过在代码实现的时候，其实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^{\prime }=W+\text{MultiHead-Attention}(W,\hat{X},\hat{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}=\alpha ||e_k|{|}^2\ast ||w_j|{|}^2+\beta$$
其中$s_{k,j}$表示文本$w_j$与目标嵌入$e_k$之间的距离，$\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_{con}+{\lambda }_i(L_{iou}+L_{dfl})$$
  其中${\lambda }_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的伪标注。

其生成流程如下：

1. 提取名词短语： n-gram
2. 伪标签(pseudo label)： 使用预训练的开放词汇检测器（GLIP），给每个图像给定的名词短语生成检测框，从而提供粗略的region-text pairs
3. 过滤(filter)： 使用预训练的CLIP来评估region-text pairs的相关性，过滤掉相关性低的注释和图像。给定图像I，图像描述T和粗略的region-text对$(B_i，t_i，c_i)$
   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=\sqrt{(c_i\ast s_i^{region})}$
   3.5. 区域级过滤(Region-level filtering)：根据文本分成不同的组，同类别使用NMS(0.5)，过滤置信度低的(0.3)
   3.6. 图像级过滤(Image-level filtering)：计算保留的所有region的平均分数，得到图像级置信度分数$S=\sqrt{(S(img)\ast S(region))}$，保留分数大于0.3的图像。
4. 使用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。