将物体识别视为下一个标记预测

我们的模型对"塞尔达传说：王国之泪"图像的前30个预测及其概率 ¹。

简介

这是论文将物体识别视为下一个标记预测的官方PyTorch实现，该论文已被CVPR 2024接收（亮点论文）。

@inproceedings{nxtp,
  title     = {{Object Recognition as Next Token Prediction}},
  author    = {Kaiyu Yue and Bor-Chun Chen and Jonas Geiping and Hengduo Li and Tom Goldstein and Ser-Nam Lim},
  booktitle = {Computer Vision and Pattern Recognition Conference (CVPR)},
  year      = {2024}
}

更新

2024年5月26日

添加ImageNet实验：查看 src/imagenet
在推理过程中可视化解码器层的注意力图：查看示例

2024年3月17日

发布在G70M上训练的最佳1.78B模型
导出onnx模型：docs/onnx-export

2024年3月3日

在本readme中添加示例，展示前20个预测结果
在评估指标中添加CLIP ViT-L/14作为文本嵌入模型（论文中的表A.8）

方法

本项目深入研究计算机视觉中的一个基本问题——物体识别，将图像转换为物体标签。

线性模型（如ResNet）和对比模型（如CLIP）需要在推理之前预定义标签，这限制了它们在实际应用中的灵活性。

我们将W扩展到整个文本空间，使用像LLaMA的32K标记嵌入这样的语言模型。我们的模型通过自回归处理以真正开放的方式预测标签。

此外，我们的一次性采样技术能够高效地进行大规模判别预测，例如前100个标签。

发布的模型有1.78B参数。将模型截断至0.77B参数仍能达到有竞争力的性能（论文中的表3），该模型在解码器中仅有一个transformer块。

示例

图像与前20个预测结果	注意力图	图像与前20个预测结果	注意力图
点击查看 ¹ `概率: 0.13949 - 图例` `概率: 0.12399 - 天空` `概率: 0.04723 - 云` `概率: 0.04642 - 游戏` `概率: 0.04500 - 截图` `概率: 0.03189 - 顶部` `概率: 0.03024 - 山` `概率: 0.02262 - 悬崖` `概率: 0.01790 - 世界` `概率: 0.01483 - Wii` `概率: 0.01440 - 视频` `概率: 0.01310 - 呼吸` `概率: 0.01087 - 泽奥` `概率: 0.00982 - 塞尔达` `概率: 0.00959 - 角色` `概率: 0.00865 - 岩石` `概率: 0.00816 - 林克` `概率: 0.00788 - 岛屿` `概率: 0.00624 - 冒险` `概率: 0.00591 - 女性`	注意力图信息 `解码器: 层 0: 头 25`	点击查看 ² `概率: 0.23237 - 火箭` `概率: 0.10435 - 发射` `概率: 0.06144 - 联盟号` `概率: 0.04314 - 太空` `概率: 0.03541 - 烟` `概率: 0.03249 - 天空` `概率: 0.01971 - 航天飞机` `概率: 0.01566 - 塔` `概率: 0.01551 - 巴黎` `概率: 0.01229 - 云` `概率: 0.01067 - 发射台` `概率: 0.01050 - 角` `概率: 0.00983 - 猎鹰` `概率: 0.00956 - 照片` `概率: 0.00834 - 升空` `概率: 0.00814 - 空气` `概率: 0.00779 - 任务` `概率: 0.00710 - 站` `概率: 0.00688 - 七月` `概率: 0.00647 - 卫星`	注意力图信息 `解码器: 层 0: 头 0`
点击查看 ³ `概率: 0.30731 - 狗` `概率: 0.13647 - 毛衣` `概率: 0.11870 - 帽子` `概率: 0.06812 - 围巾` `概率: 0.04131 - 砖` `概率: 0.03114 - 墙` `概率: 0.01796 - 衬衫` `概率: 0.01471 - 可爱` `概率: 0.01156 - 帽子` `概率: 0.00982 - 脖子` `概率: 0.00929 - 顶部` `概率: 0.00797 - 头` `概率: 0.00777 - 无檐帽` `概率: 0.00658 - 男人` `概率: 0.00588 - 坐着` `概率: 0.00582 - 外套` `概率: 0.00524 - 夹克` `概率: 0.00476 - 领子` `概率: 0.00460 - 脸` `概率: 0.00119 - 骨头`	注意力图信息 `解码器: 层 0: 头 25`	点击查看 ⁴ `概率: 0.14861 - 咖啡` `概率: 0.10409 - 商店` `概率: 0.08065 - 柜台` `概率: 0.04603 - 酒吧` `概率: 0.04055 - 餐厅` `概率: 0.03691 - 内部` `概率: 0.03468 - 区域` `概率: 0.02638 - 商店` `概率: 0.02219 - 桌子` `概率: 0.01930 - 室内` `概率: 0.01347 - 许多` `概率: 0.01156 - 食物` `概率: 0.01058 - 顾客` `概率: 0.01001 - 房间` `概率: 0.00923 - 星巴克` `概率: 0.00853 - 面包店` `概率: 0.00738 - 视图` `概率: 0.00738 - 地板` `概率: 0.00733 - 咖啡馆` `概率: 0.00633 - 架子`	注意力图信息 `解码器: 层 0: 头 8`
点击查看³ `概率: 0.47652 - 怪物` `概率: 0.09664 - 卡通` `概率: 0.03812 - 角色` `概率: 0.03724 - 群组` `概率: 0.03312 - 生物` `概率: 0.02111 - 可爱` `概率: 0.01929 - 矢量` `概率: 0.01481 - 动物` `概率: 0.00955 - 艺术` `概率: 0.00924 - 外星人` `概率: 0.00837 - 姿势` `概率: 0.00604 - 泡泡` `概率: 0.00553 - 眼睛` `概率: 0.00533 - 颜色` `概率: 0.00528 - 手` `概率: 0.00477 - 设计` `概率: 0.00474 - 壁纸` `概率: 0.00462 - 孩子` `概率: 0.00445 - 人物` `概率: 0.00445 - 家庭`	注意力图信息 `解码器: 第2层: 第7头`	点击查看³ `概率: 0.54375 - 云` `概率: 0.09932 - 词` `概率: 0.07571 - 天空` `概率: 0.03153 - 字母` `概率: 0.01862 - 索拉` `概率: 0.01380 - 标志` `概率: 0.00995 - 文本` `概率: 0.00715 - 顶部` `概率: 0.00715 - 蓝色` `概率: 0.00677 - 标题` `概率: 0.00608 - 照片` `概率: 0.00427 - 图片` `概率: 0.00288 - 索诺拉` `概率: 0.00269 - 中间` `概率: 0.00257 - 风暴` `概率: 0.00202 - 云景` `概率: 0.00190 - 太阳` `概率: 0.00189 - 艺术` `概率: 0.00156 - 翱翔` `概率: 0.00041 - 结冰的`	注意力图信息 `解码器: 第1层: 第13头`
点击查看³ `概率: 0.15317 - 建筑` `概率: 0.13619 - 波浪` `概率: 0.04782 - 房间` `概率: 0.03498 - 中间` `概率: 0.03188 - 大厅` `概率: 0.02367 - 人群` `概率: 0.02135 - 海洋` `概率: 0.02087 - 地板` `概率: 0.01867 - 世界` `概率: 0.01773 - 内部` `概率: 0.01548 - 男人` `概率: 0.01380 - 水` `概率: 0.01205 - 视图` `概率: 0.01200 - 冲浪者` `概率: 0.01109 - 照片` `概率: 0.00798 - 酒店` `概率: 0.00734 - 城市` `概率: 0.00662 - 游泳池` `概率: 0.00566 - 艺术` `概率: 0.00319 - 壁画`	注意力图信息 `解码器: 第1层: 第16头`	点击查看³ `概率: 0.25673 - 鸟` `概率: 0.21676 - 羽毛` `概率: 0.18550 - 孔雀` `概率: 0.04251 - 头部` `概率: 0.03240 - 蓝色` `概率: 0.02507 - 鸽子` `概率: 0.02183 - 尾巴` `概率: 0.01339 - 毛发` `概率: 0.01187 - 顶部` `概率: 0.00677 - 脸` `概率: 0.00631 - 相机` `概率: 0.00463 - 喙` `概率: 0.00451 - 眼睛` `概率: 0.00419 - 栅栏` `概率: 0.00370 - 坐着` `概率: 0.00333 - 栖息` `概率: 0.00330 - 照片` `概率: 0.00318 - 墙` `概率: 0.00269 - 动物` `概率: 0.00106 - 松鸦`	注意力图信息 `解码器: 第1层: 第25头`
点击查看 ⁵ `概率: 0.07247 - 平板电脑` `概率: 0.06770 - 咖啡` `概率: 0.06562 - 窗户` `概率: 0.05829 - 控制器` `概率: 0.05668 - 游戏` `概率: 0.04802 - 开关` `概率: 0.04043 - Wii` `概率: 0.03798 - 游戏机` `概率: 0.03563 - 杯子` `概率: 0.02570 - 顶部` `概率: 0.02067 - 马克杯` `概率: 0.01808 - 屏幕` `概率: 0.01344 - 视频` `概率: 0.01105 - 星星` `概率: 0.01092 - 任天堂` `概率: 0.01055 - 电脑` `概率: 0.00819 - 马里奥` `概率: 0.00815 - 遥控器` `概率: 0.00736 - 控制` `概率: 0.00393 - 窗台`	注意力图信息 `解码器: 层 0: 头 12`	点击查看 ⁶ `概率: 0.36523 - 飞机` `概率: 0.09151 - 货物` `概率: 0.07531 - 飞机` `概率: 0.05538 - 船` `概率: 0.04223 - 集装箱` `概率: 0.03105 - 水` `概率: 0.03040 - 视图` `概率: 0.02277 - 码头` `概率: 0.01685 - 港口` `概率: 0.01434 - 天空` `概率: 0.01328 - 航运` `概率: 0.00788 - 中间` `概率: 0.00751 - 机身` `概率: 0.00717 - 照片` `概率: 0.00715 - 喷气机` `概率: 0.00714 - 城市` `概率: 0.00621 - 海洋` `概率: 0.00615 - 货运` `概率: 0.00609 - 船` `概率: 0.00320 - 运输`	注意力图信息 `解码器: 层 2: 头 14`
点击查看 ⁶ `概率: 0.15236 - 糖果` `概率: 0.12271 - 毛衣` `概率: 0.11457 - 眼镜` `概率: 0.10593 - 狗` `概率: 0.08311 - 椅子` `概率: 0.07111 - 手杖` `概率: 0.04701 - 太阳镜` `概率: 0.04589 - 圣诞` `概率: 0.02361 - 服装` `概率: 0.02085 - 穿着` `概率: 0.01870 - 帽子` `概率: 0.00734 - 头部` `概率: 0.00636 - 顶部` `概率: 0.00577 - 装扮` `概率: 0.00520 - 巧克力` `概率: 0.00437 - 霍利` `概率: 0.00362 - 西装` `概率: 0.00344 - 衬衫` `概率: 0.00322 - 草莓` `概率: 0.00211 - 假发`	注意力图信息 `解码器: 层 1: 头 16`	点击查看 ⁶ `概率: 0.19960 - 客厅` `概率: 0.16291 - 房间` `概率: 0.11353 - 沙发` `概率: 0.06036 - 长沙发` `概率: 0.04741 - 地毯` `概率: 0.04704 - 咖啡` `概率: 0.03795 - 狗` `概率: 0.03659 - 墙` `概率: 0.02980 - 桌子` `概率: 0.01611 - 地板` `概率: 0.01594 - 灰色` `概率: 0.01472 - 木头` `概率: 0.01353 - 家具` `概率: 0.01314 - 植物` `概率: 0.01274 - 壁炉` `概率: 0.01161 - 枕头` `概率: 0.00941 - 椅子` `概率: 0.00512 - 家` `概率: 0.00434 - 毯子` `概率: 0.00351 - 艺术`	注意力图信息 `解码器: 层 1: 头 16`

模型

下表显示了使用前10个预测在验证集上复现的召回率结果（论文表1中的R列）。

参数量	训练组	检查点	md5	CC3M	COCO	OpenImages
1.78B	G3M	Hugging Face	`b2a69b`	0.740	0.703	0.616
1.78B	G70M	Hugging Face	`e177c7`	0.721	0.765	0.662

下载

可以从上表中的链接下载检查点。对于从Hugging Face下载，一种选择是使用git-lfs：

# 安装git lfs
git lfs install

# 在终端中下载检查点
git clone https://huggingface.co/kaiyuyue/nxtp

此外，也可以从网页浏览器的模型页面下载检查点。

推理

这里有一张图片assets/starbux.jpg用于快速测试。首先，请按照Dependencies中的说明准备环境。

要对图像进行推理，请运行

python src/infer.py \
  --ckpt-path path/to/model/checkpoint \
  --img-path assets/starbux.jpg \
  --num-labels 20

在G3M上训练的模型输出将是

前20个预测：
| 概率：0.05742 - 咖啡
| 概率：0.05525 - 餐厅
| 概率：0.04402 - 商店
| 概率：0.02528 - 房间
| 概率：0.02468 - 店铺
| 概率：0.02381 - 室内
| 概率：0.01732 - 区域
| 概率：0.01640 - 建筑
| 概率：0.01616 - 食物
| 概率：0.01408 - 酒吧
| 概率：0.01247 - 顾客
| 概率：0.01134 - 视图
| 概率：0.01059 - 地板
| 概率：0.01045 - 桌子
| 概率：0.00933 - 厨房
| 概率：0.00926 - 家
| 概率：0.00872 - 看
| 概率：0.00841 - 人
| 概率：0.00693 - 杯子
| 概率：0.00665 - 柜台

在G70M上训练的模型输出将是

前20个预测：
| 概率：0.15203 - 咖啡
| 概率：0.09728 - 商店
| 概率：0.09182 - 柜台
| 概率：0.03848 - 室内
| 概率：0.03389 - 酒吧
| 概率：0.03215 - 餐厅
| 概率：0.02440 - 桌子
| 概率：0.02245 - 店铺
| 概率：0.01950 - 区域
| 概率：0.01905 - 内部
| 概率：0.01590 - 星巴克
| 概率：0.01313 - 咖啡馆
| 概率：0.01220 - 椅子
| 概率：0.01172 - 地板
| 概率：0.01020 - 杯子
| 概率：0.00879 - 饮料
| 概率：0.00794 - 房间
| 概率：0.00746 - 顾客
| 概率：0.00635 - 木头
| 概率：0.00345 - 面包店