大型视觉编码器的自回归预训练

本仓库是 AIM 系列的所有相关内容入口，AIM 是一系列突破视觉和多模态学习边界的自回归模型：

• AIMv2: 大型视觉编码器的多模态自回归预训练 [BibTeX] [CVPR 2025（亮点）]
  Enrico Fini*, Mustafa Shukor*, Xiujun Li, Philipp Dufter, Michal Klein, David Haldimann, Sai Aitharaju, Victor Guilherme Turrisi da Costa, Louis Béthune, Zhe Gan, Alexander T Toshev, Marcin Eichner, Moin Nabi, Yinfei Yang, Joshua M. Susskind, 和 Alaaeldin El-Nouby*
• AIMv1: 大型自回归图像模型的可扩展预训练 [BibTeX][ICML 2024]
  Alaaeldin El-Nouby, Michal Klein, Shuangfei Zhai, Miguel Angel Bautista, Alexander Toshev, Vaishaal Shankar, Joshua M Susskind, Armand Joulin。

*: 技术贡献相等

如果您正在寻找原始的 AIM 模型（AIMv1），请参阅此处的 README 文件 这里。

AIMv2 概述

我们推出了使用多模态自回归目标进行预训练的 AIMv2 视觉模型系列。AIMv2 的预训练过程简单直接，易于训练且能有效扩展。AIMv2 的一些亮点包括：

1. 在大多数多模态理解基准测试中，性能优于 OAI CLIP 和 SigLIP。
2. 在开放词汇的目标检测和指代表达理解任务上，性能超越 DINOv2。
3. 展现出强大的识别能力，其中 AIMv2-3B 在使用冻结主干的情况下，在 ImageNet 上达到了 89.5% 的准确率。

AIMv2 模型图库

我们与社区共享不同容量和预训练分辨率的 AIMv2 预训练检查点：

• [AIMv2 224px 版]
• [AIMv2 336px 版]
• [AIMv2 448px 版]
• [AIMv2 原生分辨率版]
• [AIMv2 蒸馏 ViT-Large 版] (推荐用于多模态应用)
• [零样本适配的 AIMv2]

安装

请按照官方 安装说明 安装 PyTorch。之后，通过以下命令安装软件包：

pip install 'git+https://github.com/apple/ml-aim.git#subdirectory=aim-v1'
pip install 'git+https://github.com/apple/ml-aim.git#subdirectory=aim-v2'

我们还为基于 Apple 芯片的研究和实验提供 MLX 后端支持。要启用 MLX 支持，只需运行：

pip install mlx

示例

使用 PyTorch

from PIL import Image

from aim.v2.utils import load_pretrained
from aim.v1.torch.data import val_transforms

img = Image.open(...)
model = load_pretrained("aimv2-large-patch14-336", backend="torch")
transform = val_transforms(img_size=336)

inp = transform(img).unsqueeze(0)
features = model(inp)

使用 MLX

from PIL import Image
import mlx.core as mx

from aim.v2.utils import load_pretrained
from aim.v1.torch.data import val_transforms

img = Image.open(...)
model = load_pretrained("aimv2-large-patch14-336", backend="mlx")
transform = val_transforms(img_size=336)

inp = transform(img).unsqueeze(0)
inp = mx.array(inp.numpy())
features = model(inp)

使用 JAX

from PIL import Image
import jax.numpy as jnp

from aim.v2.utils import load_pretrained
from aim.v1.torch.data import val_transforms

img = Image.open(...)
model, params = load_pretrained("aimv2-large-patch14-336", backend="jax")
transform = val_transforms(img_size=336)

inp = transform(img).unsqueeze(0)
inp = jnp.array(inp)
features = model.apply({"params": params}, inp)

预训练检查点

这些预训练模型可以通过 HuggingFace Hub 访问，如下所示：

from PIL import Image
from transformers import AutoImageProcessor, AutoModel

image = Image.open(...)
processor = AutoImageProcessor.from_pretrained("apple/aimv2-large-patch14-336")
model = AutoModel.from_pretrained("apple/aimv2-large-patch14-336", trust_remote_code=True)

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

AIMv2 224px 分辨率版本

AIMv2 336px 分辨率版本

AIMv2 448px 分辨率版本

AIMv2 原生分辨率版本

我们还提供了一个经过微调的 AIMv2-L 检查点，能够处理广泛的图像分辨率和宽高比。无论图像的宽高比如何，都会将其划分为大小为 14×14 的补丁，并在经过线性投影后的输入补丁上添加 二维正弦位置编码。 该检查点支持的补丁数量范围为 [112, 4096]。

AIMv2 蒸馏版 ViT-Large

我们提供了一个从 AIMv2-3B 蒸馏而来的 AIMv2-L 检查点，它在多模态理解基准测试中表现出色。

零样本适配的 AIMv2

我们提供了经过 LiT 微调后的 AIMv2-L 视觉和文本编码器，以实现零样本识别。

引用

如果您觉得我们的工作有用，请考虑引用我们：

AIMv2 bibtex

@misc{fini2024multimodal,
    title={Multimodal Autoregressive Pre-training of Large Vision Encoders},
    author={Enrico Fini and Mustafa Shukor and Xiujun Li and Philipp Dufter and Michal Klein and David Haldimann and Sai Aitharaju and Victor Guilherme Turrisi da Costa and Louis Béthune and Zhe Gan and Alexander T Toshev and Marcin Eichner and Moin Nabi and Yinfei Yang and Joshua M. Susskind and Alaaeldin El-Nouby},
    year={2024},
    eprint={2411.14402},
    archivePrefix={arXiv},
    primaryClass={cs.CV}
}

AIMv1 bibtex

@InProceedings{pmlr-v235-el-nouby24a,
  title     = {Scalable Pre-training of Large Autoregressive Image Models},
  author    = {El-Nouby, Alaaeldin and Klein, Michal and Zhai, Shuangfei and Bautista, Miguel \'{A}ngel and Shankar, Vaishaal and Toshev, Alexander T and Susskind, Joshua M. and Joulin, Armand},
  booktitle = {Proceedings of the 41st International Conference on Machine Learning},
  pages     = {12371--12384},
  year      = {2024},
}

许可证

请在使用提供的代码和模型之前，查看仓库中的 LICENSE 文件。

ml-aim 快速上手指南

ml-aim 是苹果开源的一系列自回归视觉预训练模型（AIM），其中 AIMv2 通过多模态自回归目标进行预训练，在多项视觉和多模态理解基准测试中超越了 CLIP、SigLIP 和 DINOv2 等主流模型。本指南将帮助你快速在本地部署并使用 AIMv2 模型。

环境准备

在开始之前，请确保你的开发环境满足以下要求：

• 操作系统: Linux, macOS (支持 Apple Silicon), 或 Windows
• Python: 3.8 或更高版本
• 核心依赖:
  • PyTorch: 必须预先安装。请访问 PyTorch 官网 获取适合你环境的安装命令。
  • Pillow: 用于图像处理 (pip install pillow)
• 可选后端:
  • MLX: 如果你使用 Apple Silicon (M1/M2/M3 系列芯片) 进行推理或实验，建议安装 MLX 以获得更佳性能。
  • JAX: 支持 JAX 后端。

提示: 国内开发者若遇到 PyTorch 下载缓慢，可使用清华或中科大镜像源安装 PyTorch。

安装步骤

你可以直接通过 GitHub 源码安装 ml-aim 库。该仓库同时包含 AIMv1 和 AIMv2，需分别指定子目录安装。

1. 安装基础包 (PyTorch 后端)

pip install 'git+https://github.com/apple/ml-aim.git#subdirectory=aim-v1'
pip install 'git+https://github.com/apple/ml-aim.git#subdirectory=aim-v2'

加速建议: 如果 GitHub 连接不稳定，可尝试配置 Git 代理或使用国内代码托管平台的镜像（如有）。

2. 安装 MLX 后端 (仅限 Apple Silicon)

如果你计划在 Mac 上使用 MLX 加速：

pip install mlx

基本使用

以下示例展示如何加载预训练的 AIMv2 模型并进行图像特征提取。默认使用 PyTorch 后端。

使用 PyTorch 进行推理

这是最通用的使用方式，适用于大多数场景。

from PIL import Image
from aim.v2.utils import load_pretrained
from aim.v1.torch.data import val_transforms

# 1. 加载图像
img = Image.open("your_image.jpg") # 替换为你的图片路径

# 2. 加载预训练模型
# 模型标识符格式：aimv2-{size}-patch14-{resolution}
# 例如：aimv2-large-patch14-336
model = load_pretrained("aimv2-large-patch14-336", backend="torch")

# 3. 准备数据预处理
# 注意：transform 的 img_size 需与模型分辨率一致 (此处为 336)
transform = val_transforms(img_size=336)

# 4. 执行推理
inp = transform(img).unsqueeze(0)
features = model(inp)

print(f"输出特征形状：{features.shape}")

常用模型推荐

根据官方评测，针对不同需求推荐以下模型：

注意: 模型权重首次运行时会自动从 Hugging Face 下载。国内用户若下载缓慢，建议配置 HF_ENDPOINT 环境变量使用镜像站：

export HF_ENDPOINT=https://hf-mirror.com