yolos-small

YOLOS-small项目介绍

YOLOS-small是一个用于目标检测的小型视觉模型,它基于视觉Transformer架构,在COCO 2017数据集上进行了微调。这个模型由Fang等人在论文《You Only Look at One Sequence: Rethinking Transformer in Vision through Object Detection》中提出,并首次在GitHub上发布。

模型特点

YOLOS-small模型具有以下几个特点:

1. 简单而高效:尽管结构简单,但基础版YOLOS模型在COCO验证集上可以达到42 AP的性能,与DETR等更复杂的框架相当。

2. 基于Transformer:它采用了Vision Transformer(ViT)的架构,但使用DETR的损失函数进行训练。

3. 双向匹配损失:模型使用匈牙利算法对预测结果和真实标注进行一一匹配,然后用交叉熵和边界框损失进行优化。

4. 预训练+微调:模型首先在ImageNet-1k上预训练200轮,然后在COCO 2017上微调150轮。

使用方法

YOLOS-small可以直接用于目标检测任务。使用时,需要先加载特征提取器和模型:

from transformers import YolosFeatureExtractor, YolosForObjectDetection

feature_extractor = YolosFeatureExtractor.from_pretrained('hustvl/yolos-small')
model = YolosForObjectDetection.from_pretrained('hustvl/yolos-small')

然后对输入图像进行处理并传入模型:

inputs = feature_extractor(images=image, return_tensors="pt")
outputs = model(**inputs)

模型会输出类别logits和边界框预测结果。

性能表现

在COCO 2017验证集上,YOLOS-small达到了36.1的AP(平均精度)。虽然不及大型号的YOLOS模型,但考虑到其小巧的模型尺寸,这个性能已经相当不错。

应用场景

YOLOS-small适用于各种目标检测场景,如自动驾驶、安防监控、图像分析等。由于模型较小,它特别适合于资源受限的设备或需要实时处理的应用。

局限性

作为一个小型模型,YOLOS-small在复杂场景或小目标检测上的表现可能不如大型模型。此外,目前只支持PyTorch框架,使用其他深度学习框架的用户可能需要进行转换。

总的来说,YOLOS-small为目标检测任务提供了一个简单高效的解决方案,特别适合于对模型大小和推理速度有要求的应用场景。