AIM：自回归图像模型

Alaaeldin El-Nouby, Michal Klein, Shuangfei Zhai, Miguel Angel Bautista, Alexander Toshev, Vaishaal Shankar, Joshua M Susskind, 和 Armand Joulin

将发表于ICML 2024

[论文] [引用]

本软件项目配合研究论文大规模自回归图像模型的可扩展预训练。

我们推出了AIM，这是一系列使用自回归生成目标预训练的视觉模型。我们展示了图像特征的自回归预训练具有与其文本对应物（即大型语言模型）相似的扩展性。具体而言，我们强调两个发现：

模型容量可以轻松扩展到数十亿参数，并且
AIM能有效利用大量未经整理的图像数据。

安装

请按照官方安装说明安装PyTorch。之后，按如下方式安装软件包：

pip install git+https://git@github.com/apple/ml-aim.git

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

pip install mlx

使用

以下是PyTorch中的使用示例：

from PIL import Image

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

img = Image.open(...)
model = load_pretrained("aim-600M-2B-imgs", backend="torch")
transform = val_transforms()

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

from PIL import Image
import mlx.core as mx

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

img = Image.open(...)
model = load_pretrained("aim-600M-2B-imgs", backend="mlx")
transform = val_transforms()

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

</details> <details> <summary>以及在<a href="https://jax.readthedocs.io/">JAX</a>中的使用</summary>

from PIL import Image
import jax.numpy as jnp

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

img = Image.open(...)
model, params = load_pretrained("aim-600M-2B-imgs", backend="jax")
transform = val_transforms()

inp = transform(img).unsqueeze(0)
inp = jnp.array(inp)
(logits, features), _ = model.apply(params, inp, mutable=['batch_stats'])

</details>

预训练检查点

可以通过PyTorch Hub访问预训练模型：

import torch

aim_600m = torch.hub.load("apple/ml-aim", "aim_600M")
aim_1b   = torch.hub.load("apple/ml-aim", "aim_1B")
aim_3b   = torch.hub.load("apple/ml-aim", "aim_3B")
aim_7b   = torch.hub.load("apple/ml-aim", "aim_7B")

或通过HuggingFace Hub访问：

from aim.torch.models import AIMForImageClassification

aim_600m = AIMForImageClassification.from_pretrained("apple/aim-600M")
aim_1b   = AIMForImageClassification.from_pretrained("apple/aim-1B")
aim_3b   = AIMForImageClassification.from_pretrained("apple/aim-3B")
aim_7b   = AIMForImageClassification.from_pretrained("apple/aim-7B")

预训练骨干网络

下表包含我们论文中使用的预训练骨干网络。

预训练注意力头

下表包含在ImageNet-1k验证集上的分类结果。

复现IN-1k分类结果

以下命令可以在ImageNet-1k验证集上复现注意力探针结果。我们使用1个节点和8个GPU进行评估：

torchrun --standalone --nnodes=1 --nproc-per-node=8 main_attnprobe.py \
  --model=aim-7B \
  --batch-size=64 \
  --data-path=/path/to/imagenet \
  --probe-layers=best \
  --backbone-ckpt-path=/path/to/backbone_ckpt.pth \
  --head-ckpt-path=/path/to/head_ckpt.pth

默认情况下，我们探测提供最佳性能的中间6层的特征。要更改此设置，只需传递--probe-layers=last。

引用

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

@article{el2024scalable,
  title={大型自回归图像模型的可扩展预训练},
  author={El-Nouby, Alaaeldin 和 Klein, Michal 和 Zhai, Shuangfei 和 Bautista, Miguel Angel 和 Toshev, Alexander 和 Shankar, Vaishaal 和 Susskind, Joshua M 和 Joulin, Armand},
  journal={国际机器学习会议},
  year={2024}
}