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

news2024/12/23 13:26:02

在这里插入图片描述

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$ 是来自不同层的图像特征，W是text embedding，计算特征图与text embedding的相似度矩阵，取最大值在sigmoid后作为权重加权原特征图
$X_l^{\prime}=X_l \cdot \delta\left(\max _{j \in\{1 . . C\}}\left(X_l W_j^{\top}\right)\right)^{\top}$
不过在代码实现的时候，其实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*||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_{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的伪标注。
在这里插入图片描述
其生成流程如下：

提取名词短语： n-gram
伪标签(pseudo label)： 使用预训练的开放词汇检测器（GLIP），给每个图像给定的名词短语生成检测框，从而提供粗略的region-text pairs
过滤(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 * s^{region}_i)}$
3.5. 区域级过滤(Region-level filtering)：根据文本分成不同的组，同类别使用NMS(0.5)，过滤置信度低的(0.3)
3.6. 图像级过滤(Image-level filtering)：计算保留的所有region的平均分数，得到图像级置信度分数 $\sqrt{(S(img) * S(region))}$ ，保留分数大于0.3的图像。
使用NMS过滤冗余的检测框。

训练集数据情况
测试集数据情况： 在LVIS进行zero-shot evaluation
- 包含1203个类别，远超训练集的类别
- 主要在LVIS minival上测试固定AP
- 最大预测数量设置为1000
  
  总结：可以看出检测效果非常Amazing，而且更关键的是速度非常快！ 其速度快的原因在于测试的时候，可以直接将text embeding先算出来，让后作为YOLO-World的参数来进行前向推理。直接砍掉了CLIP在推理阶段的耗时。这个思路非常nice！