简而言之： 我们研究了在中等规模的 ImageNet-1k 数据集上预训练的普通 ViT 模型，在更具挑战性的 COCO 目标检测基准上的迁移能力。

:man_technologist: 该项目目前处于积极开发中 :woman_technologist: :

• 2022年5月4日: :eyes:YOLOS 现已在 🤗HuggingFace Transformers 中发布！

• 2022年4月8日: 如果你喜欢 YOLOS，你可能也会喜欢 MIMDet（论文 / 代码与模型)！MIMDet 能够高效且有效地将基于掩码图像建模（MIM）预训练的普通视觉 Transformer（ViT）适配到高性能目标检测任务中（使用 ViT-Base 和 Mask R-CNN 在 COCO 上可达到 51.5 的 box AP 和 46.0 的 mask AP）。

• 2021年10月28日: YOLOS 针对 NeurIPS 2021 的最终定稿版本进行了更新。我们加入了 MoCo-v3 自监督预训练的结果，研究了分离 [Det] tokens 的影响，并新增了一个讨论部分。

• 2021年9月29日: YOLOS 已被 NeurIPS 2021 接收！

• 2021年6月22日: 我们在 arXiv 上更新了我们的 论文，增加了关于位置嵌入的讨论和更多可视化内容，快来看看吧！

• 2021年6月9日: 我们添加了一个 notebook，用于可视化最后一层不同头上的 [Det] tokens 的自注意力图，欢迎查看！

你只看一个序列：通过目标检测重新思考视觉中的 Transformer

作者：Yuxin Fang¹ *, Bencheng Liao¹ *, Xinggang Wang^{1 :email:}, Jiemin Fang^{2, 1}, Jiyang Qi¹, Rui Wu³, Jianwei Niu³, Wenyu Liu¹。

¹ 华中科技大学电子信息学院, ² 华中科技大学人工智能研究院, ³ 地平线机器人。

(*) 共同第一作者，(^:email:) 通讯作者。

arXiv 技术报告 (arXiv 2106.00666)

你只看一个序列（YOLOS）

YOLOS 的示意图

亮点

YOLOS 直接继承自 ViT（DeiT），它并非旨在成为又一个高性能的目标检测器，而是为了揭示 Transformer 从图像识别到目标检测的通用性和迁移性。 具体来说，我们的主要贡献总结如下：

• 我们仅使用中等规模的 ImageNet-1k 作为唯一的预训练数据集，并证明普通的 ViT（DeiT）只需进行最少的修改，就能成功迁移到具有挑战性的目标检测任务中，以“只看一个序列”的方式（YOLOS），并取得与现有方法相当的 COCO 结果。

• 我们展示了可以通过纯序列到序列的方式完成 2D 目标检测，即以固定大小、不重叠的图像块序列作为输入。在现有的目标检测器中，YOLOS 使用的 2D 归纳偏置最少。此外，YOLOS 还可以在任何维度的空间中执行目标检测，而无需了解精确的空间结构或几何形状。

• 对于 ViT（DeiT），我们发现其目标检测结果对预训练方案非常敏感，且检测性能远未达到饱和状态。因此，提出的 YOLOS 可以作为一个具有挑战性的基准任务，用于评估 ViT（DeiT）的不同预训练策略。

• 我们还通过目标检测任务的迁移，探讨了当前视觉领域中流行的预训练方案和模型缩放策略的影响及局限性。

结果

注释:

• 百度网盘的提取码为yolo。
• FB代表由 DeiT 提供的模型权重（论文, 代码）。感谢他们的优秀工作。
• 我们将在未来更新其他模型，请持续关注 :)

要求

本代码库使用 Python 3.6、PyTorch 1.5+ 和 torchvision 0.6+ 开发：

conda install -c pytorch pytorch torchvision

安装 pycocotools（用于在 COCO 数据集上进行评估）和 scipy（用于训练）：

conda install cython scipy
pip install -U 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'

数据准备

从 http://cocodataset.org 下载并解压 COCO 2017 训练集和验证集的图像及其标注文件。我们期望目录结构如下：

path/to/coco/
  annotations/  # 标注 JSON 文件
  train2017/    # 训练图像
  val2017/      # 验证图像

训练

在 COCO 数据集上进行微调之前，你需要将 ImageNet 预训练模型下载到 /path/to/YOLOS/ 目录下。

python -m torch.distributed.launch \
    --nproc_per_node=8 \
    --use_env main.py \
    --coco_path /path/to/coco
    --batch_size 2 \
    --lr 5e-5 \
    --epochs 300 \
    --backbone_name tiny \
    --pre_trained /path/to/deit-tiny.pth\
    --eval_size 512 \
    --init_pe_size 800 1333 \
    --output_dir /output/path/box_model

python -m torch.distributed.launch 
    --nproc_per_node=8 
    --use_env main.py 
    --coco_path /path/to/coco
    --batch_size 1 
    --lr 2.5e-5 
    --epochs 150 
    --backbone_name small 
    --pre_trained /path/to/deit-small-200epoch.pth
    --eval_size 800 
    --init_pe_size 512 864 
    --mid_pe_size 512 864 
    --output_dir /output/path/box_model

python -m torch.distributed.launch \
    --nproc_per_node=8 \
    --use_env main.py \
    --coco_path /path/to/coco
    --batch_size 1 \
    --lr 2.5e-5 \
    --epochs 150 \
    --backbone_name small \
    --pre_trained /path/to/deit-small-300epoch.pth\
    --eval_size 800 \
    --init_pe_size 512 864 \
    --mid_pe_size 512 864 \
    --output_dir /output/path/box_model

python -m torch.distributed.launch \
    --nproc_per_node=8 \
    --use_env main.py \
    --coco_path /path/to/coco
    --batch_size 1 \
    --lr 2.5e-5 \
    --epochs 150 \
    --backbone_name small_dWr \
    --pre_trained /path/to/deit-small-dWr-scale.pth\
    --eval_size 800 \
    --init_pe_size 512 864 \
    --mid_pe_size 512 864 \
    --output_dir /output/path/box_model

python -m torch.distributed.launch \
    --nproc_per_node=8 \
    --use_env main.py \
    --coco_path /path/to/coco
    --batch_size 1 \
    --lr 2.5e-5 \
    --epochs 150 \
    --backbone_name base \
    --pre_trained /path/to/deit-base.pth\
    --eval_size 800 \
    --init_pe_size 800 1344 \
    --mid_pe_size 800 1344 \
    --output_dir /output/path/box_model

评估

要在 COCO 数据集上评估 YOLOS-Ti 模型，请运行以下命令：

python -m torch.distributed.launch --nproc_per_node=8 --use_env main.py --coco_path /path/to/coco --batch_size 2 --backbone_name tiny --eval --eval_size 512 --init_pe_size 800 1333 --resume /path/to/YOLOS-Ti

要在 COCO 数据集上评估 YOLOS-S 模型，请运行以下命令：

python -m torch.distributed.launch --nproc_per_node=8 --use_env main.py --coco_path /path/to/coco --batch_size 1 --backbone_name small --eval --eval_size 800 --init_pe_size 512 864 --mid_pe_size 512 864 --resume /path/to/YOLOS-S

要在 COCO 数据集上评估 YOLOS-S (dWr) 模型，请运行以下命令：

python -m torch.distributed.launch --nproc_per_node=8 --use_env main.py --coco_path /path/to/coco --batch_size 1 --backbone_name small_dWr --eval --eval_size 800 --init_pe_size 512 864 --mid_pe_size 512 864 --resume /path/to/YOLOS-S(dWr)

要在 COCO 数据集上评估 YOLOS-B 模型，请运行以下命令：

python -m torch.distributed.launch --nproc_per_node=8 --use_env main.py --coco_path /path/to/coco --batch_size 1 --backbone_name base --eval --eval_size 800 --init_pe_size 800 1344 --mid_pe_size 800 1344 --resume /path/to/YOLOS-B

可视化

• 可视化边界框预测及目标类别分布：

1. 要获得论文中的可视化效果，您需要在 COCO 数据集上微调过的 YOLOS 模型。运行以下命令以获取 COCO 验证集上的 100 个 Det-Token 预测结果，随后会生成 /path/to/YOLOS/visualization/modelname-eval-800-eval-pred.json 文件：

python cocoval_predjson_generation.py --coco_path /path/to/coco --batch_size 1 --backbone_name small --eval --eval_size 800 --init_pe_size 512 864 --mid_pe_size 512 864 --resume /path/to/yolos-s-model.pth --output_dir ./visualization

2. 要获取 COCO 验证集所有图像的全部真实标注类别，运行以下命令以生成 /path/to/YOLOS/visualization/coco-valsplit-cls-dist.json 文件：

python cocoval_gtclsjson_generation.py --coco_path /path/to/coco --batch_size 1 --output_dir ./visualization

3. 要可视化 Det-Token 的边界框和类别的分布，运行以下命令以在 /path/to/YOLOS/visualization/ 目录下生成 .png 文件：

 python visualize_dettoken_dist.py --visjson /path/to/YOLOS/visualization/modelname-eval-800-eval-pred.json --cococlsjson /path/to/YOLOS/visualization/coco-valsplit-cls-dist.json

• 使用 VisualizeAttention.ipynb 可视化最后一层不同头上的 [Det] Token 自注意力：

致谢 :heart:

本项目基于 DETR（论文，代码)、DeiT（论文，代码)、DINO（论文，代码）以及 timm。感谢这些优秀的工作。

引用

如果您在研究中发现我们的论文和代码有所帮助，请考虑给个项目点赞 :star: 并引用 :pencil: 我们：

@article{YOLOS,
  title={You Only Look at One Sequence: Rethinking Transformer in Vision through Object Detection},
  author={Fang, Yuxin and Liao, Bencheng and Wang, Xinggang and Fang, Jiemin and Qi, Jiyang and Wu, Rui and Niu, Jianwei and Liu, Wenyu},
  journal={arXiv preprint arXiv:2106.00666},
  year={2021}
}

YOLOS 快速上手指南

YOLOS (You Only Look at One Sequence) 是一个基于纯 Transformer 架构的目标检测模型。它直接复用预训练的 ViT/DeiT 权重，仅需极少修改即可在 COCO 数据集上实现具有竞争力的检测效果，旨在探索 Transformer 从图像识别到目标检测的迁移能力。

1. 环境准备

本项目基于 Python 3.6+ 开发，依赖 PyTorch 1.5+ 和 torchvision 0.6+。

系统要求：

• Python >= 3.6
• PyTorch >= 1.5
• torchvision >= 0.6
• CUDA 环境（推荐用于训练和推理）

前置依赖安装： 建议优先使用国内镜像源加速安装。

# 安装 PyTorch 和 torchvision (推荐使用清华或阿里镜像)
pip install torch torchvision -i https://pypi.tuna.tsinghua.edu.cn/simple

# 安装其他必要依赖
pip install cython scipy -i https://pypi.tuna.tsinghua.edu.cn/simple

# 安装 pycocotools (COCO 评估必备)
pip install -U 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI' -i https://pypi.tuna.tsinghua.edu.cn/simple

2. 数据准备

下载 COCO 2017 数据集（训练集、验证集及标注文件），并解压至指定目录。目录结构需如下所示：

path/to/coco/
  annotations/  # 包含 json 标注文件
  train2017/    # 训练图片
  val2017/      # 验证图片

提示：请提前从 ImageNet 预训练好的 DeiT/ViT 权重（如 deit-tiny.pth, deit-small.pth 等），并将其放置在项目目录下备用。

3. 基本使用

模型训练 (Fine-tuning)

在 COCO 数据集上进行微调前，请确保已下载对应的 ImageNet 预训练权重。以下以训练 YOLOS-Ti (Tiny) 模型为例，使用 8 张 GPU 进行分布式训练：

python -m torch.distributed.launch \
    --nproc_per_node=8 \
    --use_env main.py \
    --coco_path /path/to/coco \
    --batch_size 2 \
    --lr 5e-5 \
    --epochs 300 \
    --backbone_name tiny \
    --pre_trained /path/to/deit-tiny.pth \
    --eval_size 512 \
    --init_pe_size 800 1333 \
    --output_dir /output/path/box_model

注：请将 /path/to/coco 替换为实际数据路径，/path/to/deit-tiny.pth 替换为预训练权重路径。其他型号（如 small, base）只需调整 --backbone_name、--pre_trained 及相关尺寸参数。

模型评估 (Evaluation)

使用训练好的检查点进行评估。以下以评估 YOLOS-Ti 为例：

python -m torch.distributed.launch \
    --nproc_per_node=8 \
    --use_env main.py \
    --coco_path /path/to/coco \
    --batch_size 2 \
    --backbone_name tiny \
    --eval \
    --eval_size 512 \
    --init_pe_size 800 1333 \
    --resume /path/to/YOLOS-Ti_checkpoint.pth

注：--resume 后接具体的模型权重文件路径。

快速体验 (Hugging Face)

YOLOS 已集成至 Hugging Face Transformers 库，可通过以下代码快速加载预训练模型进行推理（无需配置完整训练环境）：

from transformers import YolosForObjectDetection, YolosImageProcessor
import torch
from PIL import Image
import requests

# 加载处理器和模型
processor = YolosImageProcessor.from_pretrained("hustvl/yolos-base")
model = YolosForObjectDetection.from_pretrained("hustvl/yolos-base")

# 准备图像
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
image = Image.open(requests.get(url, stream=True).raw)

# 推理
inputs = processor(images=image, return_tensors="pt")
outputs = model(**inputs)

# 后处理获取结果
target_sizes = torch.tensor([image.size[::-1]])
results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]

for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
    print(f"Detected {model.config.id2label[label.item()]} with confidence {score.item():.2f} at box {box.tolist()}")