无向量量化的自回归图像生成 官方 PyTorch 实现

这是论文无向量量化的自回归图像生成的 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 模型
💥 用于运行各种预训练 MAR 模型的独立 Colab 笔记本
🛸 使用 PyTorch DDP 的 MAR+DiffLoss 训练和评估脚本

准备工作

数据集

下载 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	参数数量
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笔记本运行我们的交互式可视化演示！

训练

默认设置的脚本（MAR-L，DiffLoss MLP有3个块和1024通道宽度，400个epochs）：

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个epochs）。
（可选）要使用缓存的VAE潜在表示进行训练，请在参数中添加--use_cached --cached_path ${CACHED_PATH}。使用缓存潜在表示的训练时间在16个H100 GPU上使用--batch_size 128约为1天11小时（比不使用缓存快近2倍）。
（可选）要在训练期间通过使用梯度检查点来节省GPU内存（感谢@Jiawei-Yang），请在参数中添加--grad_checkpointing。请注意，这可能会略微降低训练速度。

在8个80G A100 GPU上测试：

MAR-L，diffloss_d 8，diffloss_w 1280：从约62GB降至约32GB，训练时间从0.78秒/迭代增加到0.94秒/迭代。 MAR-H，diffloss_d 12，diffloss_w 1536：从内存不足（OOM）到约44GB，训练时间约为1.56秒/迭代。全局批量大小为8x64=512

评估（ImageNet 256x256）

使用无分类器引导评估MAR-B（DiffLoss MLP有6个块和1024通道宽度，800个epochs）：

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（DiffLoss MLP有8个块和1280通道宽度，800个epochs）：

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（DiffLoss MLP有12个块和1536通道宽度，800个epochs）：

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