无需向量量化自回归图像生成
_{官方 PyTorch 实现}

   

这是论文《无需向量量化的自回归图像生成》（NeurIPS 2024 Spotlight Presentation）的 PyTorch/GPU 实现：

@article{li2024autoregressive,
  title={Autoregressive Image Generation without Vector Quantization},
  author={Li, Tianhong and Tian, Yonglong and Li, He and Deng, Mingyang and He, Kaiming},
  journal={arXiv preprint arXiv:2406.11838},
  year={2024}
}

本仓库包含：

• 🪐 MAR 和 DiffLoss 的简单 PyTorch 实现
• ⚡️ 在 ImageNet 256x256 数据集上训练好的类别条件 MAR 模型
• 💥 一个自包含的 Colab 笔记本，用于运行各种预训练的 MAR 模型
• 🛸 使用 PyTorch DDP 的 MAR+DiffLoss 训练与评估脚本
• 🎉 此外，还可以查看我们的 Hugging Face 模型卡片 和 Gradio 演示（感谢 @jadechoghari）。

准备工作

数据集

下载 ImageNet 数据集，并将其放置在你的 IMAGENET_PATH 目录下。

安装

克隆代码：

git clone https://github.com/LTH14/mar.git
cd mar

可以创建并激活一个名为 mar 的合适 conda 环境：

conda env create -f environment.yaml
conda activate mar

下载预训练的 VAE 和 MAR 模型：

python util/download.py

为方便起见，我们还在此直接提供预训练的 MAR 模型下载：

MAR 模型	FID-50K	Inception Score	#params
MAR-B	2.31	281.7	208M
MAR-L	1.78	296.0	479M
MAR-H	1.55	303.7	943M

（可选）缓存 VAE 隐变量

鉴于我们的数据增强仅包括简单的中心裁剪和随机翻转，VAE 隐变量可以预先计算并保存到 CACHED_PATH，以节省 MAR 训练过程中的计算量：

torchrun --nproc_per_node=8 --nnodes=1 --node_rank=0 \
main_cache.py \
--img_size 256 --vae_path pretrained_models/vae/kl16.ckpt --vae_embed_dim 16 \
--batch_size 128 \
--data_path ${IMAGENET_PATH} --cached_path ${CACHED_PATH}

使用方法

演示

使用 Colab 笔记本运行我们的交互式可视化 演示！

本地 Gradio 应用程序

python demo/gradio_app.py 

训练

默认设置的脚本（MAR-L，具有 3 个模块、宽度为 1024 通道的 DiffLoss MLP，训练 400 个 epoch）：

torchrun --nproc_per_node=8 --nnodes=4 --node_rank=${NODE_RANK} --master_addr=${MASTER_ADDR} --master_port=${MASTER_PORT} \
main_mar.py \
--img_size 256 --vae_path pretrained_models/vae/kl16.ckpt --vae_embed_dim 16 --vae_stride 16 --patch_size 1 \
--model mar_large --diffloss_d 3 --diffloss_w 1024 \
--epochs 400 --warmup_epochs 100 --batch_size 64 --blr 1.0e-4 --diffusion_batch_mul 4 \
--output_dir ${OUTPUT_DIR} --resume ${OUTPUT_DIR} \
--data_path ${IMAGENET_PATH}

• 在 32 张 H100 GPU 上，使用 --batch_size 64 时，训练时间约为 1 天 7 小时。
• 添加 --online_eval 可以在训练过程中评估 FID（每 40 个 epoch 评估一次）。
• （可选）若要使用缓存的 VAE 隐变量进行训练，可在参数中添加 --use_cached --cached_path ${CACHED_PATH}。使用缓存隐变量时，在 16 张 H100 GPU 上，使用 --batch_size 128 的训练时间约为 1 天 11 小时（几乎比不使用缓存快 2 倍）。
• （可选）为了在训练过程中通过梯度检查点节省显存（感谢 @Jiawei-Yang），可在参数中添加 --grad_checkpointing。请注意，这可能会略微降低训练速度。

评估（ImageNet 256x256）

使用无分类器指导评估 MAR-B（具有 6 个模块、宽度为 1024 通道的 DiffLoss MLP，训练 800 个 epoch）：

torchrun --nproc_per_node=8 --nnodes=1 --node_rank=0 \
main_mar.py \
--model mar_base --diffloss_d 6 --diffloss_w 1024 \
--eval_bsz 256 --num_images 50000 \
--num_iter 256 --num_sampling_steps 100 --cfg 2.9 --cfg_schedule linear --temperature 1.0 \
--output_dir pretrained_models/mar/mar_base \
--resume pretrained_models/mar/mar_base \
--data_path ${IMAGENET_PATH} --evaluate

使用无分类器指导评估 MAR-L（具有 8 个模块、宽度为 1280 通道的 DiffLoss MLP，训练 800 个 epoch）：

torchrun --nproc_per_node=8 --nnodes=1 --node_rank=0 \
main_mar.py \
--model mar_large --diffloss_d 8 --diffloss_w 1280 \
--eval_bsz 256 --num_images 50000 \
--num_iter 256 --num_sampling_steps 100 --cfg 3.0 --cfg_schedule linear --temperature 1.0 \
--output_dir pretrained_models/mar/mar_large \
--resume pretrained_models/mar/mar_large \
--data_path ${IMAGENET_PATH} --evaluate

使用无分类器指导评估 MAR-H（具有 12 个模块、宽度为 1536 通道的 DiffLoss MLP，训练 800 个 epoch）：

torchrun --nproc_per_node=8 --nnodes=1 --node_rank=0 \
main_mar.py \
--model mar_huge --diffloss_d 12 --diffloss_w 1536 \
--eval_bsz 128 --num_images 50000 \
--num_iter 256 --num_sampling_steps 100 --cfg 3.2 --cfg_schedule linear --temperature 1.0 \
--output_dir pretrained_models/mar/mar_huge \
--resume pretrained_models/mar/mar_huge \
--data_path ${IMAGENET_PATH} --evaluate

• 设置 --cfg 1.0 --temperature 0.95 可以在没有无分类器指导的情况下进行评估。
• 通过减少自回归迭代次数（例如 --num_iter 64），可以显著提高生成速度。

致谢

我们感谢邓聪悦和陈欣蕾的有益讨论。同时，我们也感谢 Google TPU 研究云（TRC）为我们提供 TPU 使用权限，以及 Google 云平台对 GPU 资源的支持。

本仓库中的大量代码基于 MAE、MAGE 和 DiT。

联系方式

如果您有任何问题，请随时通过电子邮件（tianhong@mit.edu）与我联系。祝您使用愉快！

MAR 快速上手指南

MAR (Autoregressive Image Generation without Vector Quantization) 是一个基于 PyTorch 的自回归图像生成模型，无需向量量化即可在 ImageNet 256x256 数据集上实现高质量的图像生成。

环境准备

• 系统要求: Linux 操作系统，配备 NVIDIA GPU（推荐 H100/A100 等高性能显卡以进行训练，推理可使用消费级显卡）。
• 前置依赖:
  • Python 3.8+
  • Conda (推荐用于环境管理)
  • Git
  • CUDA Toolkit (需与 PyTorch 版本匹配)
• 数据集: 需自行下载 ImageNet 数据集，并记录其路径为 IMAGENET_PATH。

注意：国内用户若下载官方源较慢，可尝试使用学术资源镜像或手动下载模型文件后放入指定目录。

安装步骤

1. 克隆代码仓库

   git clone https://github.com/LTH14/mar.git
   cd mar
   

2. 创建并激活 Conda 环境 项目提供了 environment.yaml 配置文件，可一键安装所需依赖：

   conda env create -f environment.yaml
   conda activate mar
   

3. 下载预训练模型 运行脚本自动下载预训练的 VAE 和 MAR 模型：

   python util/download.py
   

   若自动下载失败，可手动从以下地址下载并放入 pretrained_models 目录：

   • MAR-B (208M)
   • MAR-L (479M)
   • MAR-H (943M)

基本使用

方式一：在线体验 (推荐新手)

无需本地配置，直接在 Google Colab 中运行交互式 Demo：

• 打开 Colab Notebook

方式二：本地运行 Gradio 演示

在本地启动一个 Web 界面进行图像生成测试：

python demo/gradio_app.py 

启动后，终端会显示访问地址（通常为 http://127.0.0.1:7860），在浏览器打开即可选择类别生成图像。

方式三：命令行推理 (评估)

使用预训练的 MAR-L 模型生成 50,000 张图像并评估 FID（需配置好 IMAGENET_PATH）：

torchrun --nproc_per_node=8 --nnodes=1 --node_rank=0 \
main_mar.py \
--model mar_large --diffloss_d 8 --diffloss_w 1280 \
--eval_bsz 256 --num_images 50000 \
--num_iter 256 --num_sampling_steps 100 --cfg 3.0 --cfg_schedule linear --temperature 1.0 \
--output_dir pretrained_models/mar/mar_large \
--resume pretrained_models/mar/mar_large \
--data_path ${IMAGENET_PATH} --evaluate

• 加速提示: 若仅需快速查看生成效果，可减少自回归迭代次数，例如添加参数 --num_iter 64。
• 无分类器引导: 若不使用 classifier-free guidance，设置 --cfg 1.0 --temperature 0.95。