U-ViT
_{All are Worth Words: A ViT Backbone for Diffusion Models（CVPR 2023）的官方 PyTorch 实现}

💡使用 U-ViT 的项目：

• UniDiffuser，一个基于 1B 参数 U-ViT 的多模态大规模扩散模型，已开源
• DPT，代码，演示，是一种仅使用 1 个标签/类训练的条件扩散模型，在 ImageNet 上实现了 SOTA 的自监督学习生成与分类效果

视觉Transformer（ViT）在各类视觉任务中展现出巨大潜力，而基于卷积神经网络（CNN）的 U-Net 仍然在扩散模型领域占据主导地位。
我们设计了一种简单通用的基于 ViT 的架构（命名为 U-ViT），用于扩散模型的图像生成。
U-ViT 的特点在于将所有输入，包括时间、条件信息以及噪声图像块，都视为 token，并在浅层和深层之间采用长跳跃连接。
我们在无条件和类别条件图像生成以及文本到图像生成任务中对 U-ViT 进行了评估，结果表明，U-ViT 在性能上不逊于甚至优于同等规模的基于 CNN 的 U-Net。
尤其值得注意的是，使用 U-ViT 的潜在空间扩散模型在 ImageNet 256x256 上的类别条件图像生成任务中取得了创纪录的 FID 分数 2.29，在 MS-COCO 数据集上的文本到图像生成任务中则达到了 5.48，且这些方法在生成模型训练过程中并未使用大型外部数据集。

我们的研究结果表明，对于基于扩散的图像建模而言，长跳跃连接至关重要，而 CNN 基础 U-Net 中的下采样和上采样操作并非总是必要的。我们相信，U-ViT 可以为未来扩散模型骨干网络的研究提供新的思路，并有助于大规模跨模态数据集上的生成建模任务。

---

本代码库实现了扩散模型中的 Transformer 骨干 📌U-ViT📌，如 论文 所介绍。U-ViT 将所有输入视为 token，并采用长跳跃连接。长跳跃连接极大地提升了性能和收敛速度。

💡本代码库包含：

• 经优化注意力计算的 U-ViT 实现
• 在常见图像生成基准数据集上预训练的 U-ViT 模型（CIFAR10、CelebA 64x64、ImageNet 64x64、ImageNet 256x256、ImageNet 512x512）
• 针对 像素空间扩散模型、潜在空间扩散模型 和 文本到图像扩散模型 的高效训练脚本
• 针对 像素空间扩散模型、潜在空间扩散模型 和 文本到图像扩散模型 的高效评估脚本
• 一个 Colab 笔记本演示，用于从 U-ViT 生成 ImageNet 样本（FID=2.29）

💡本代码库支持多种用于高效训练和采样的扩散模型技术：

• 使用 huggingface accelerate 库进行混合精度训练（🥰自动开启）
• 使用 facebook xformers 库实现高效的注意力计算（需额外安装）
• 梯度检查点技术，可减少约 65% 的显存占用（🥰自动开启）
• 结合这些技术，我们能够在仅配备 2 张 A100 显卡的情况下，以 1024 的大批次训练最大规模的 U-ViT-H 模型，分辨率高达 256x256 和 512x512❗

使用一张 A100 显卡，在 ImageNet 256x256 上以 128 的批次训练 U-ViT-H/2 时的训练速度和显存消耗：

混合精度训练	xformers	梯度检查点	训练速度	显存
❌	❌	❌	-	显存不足
✔	❌	❌	0.97 步/秒	78852 MB
✔	✔	❌	1.14 步/秒	54324 MB
✔	✔	✔	0.87 步/秒	18858 MB

依赖项

pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu116  # 安装 torch-1.13.1
pip install accelerate==0.12.0 absl-py ml_collections einops wandb ftfy==6.1.1 transformers==4.23.1

# xformers 是可选的，但能显著加速注意力计算。
pip install -U xformers
pip install -U --pre triton

• 本仓库基于 timm==0.3.2，该版本需要应用一个 修复，才能与 PyTorch 1.8.1+ 兼容。（其他版本或许也能工作，但我尚未测试过。）
• 我们强烈建议安装 xformers，它能极大提升 训练和推理 过程中的注意力计算速度。

预训练模型

模型	FID	训练迭代次数	批量大小
CIFAR10 (U-ViT-S/2)	3.11	50万	128
CelebA 64x64 (U-ViT-S/4)	2.87	50万	128
ImageNet 64x64 (U-ViT-M/4)	5.85	30万	1024
ImageNet 64x64 (U-ViT-L/4)	4.26	30万	1024
ImageNet 256x256 (U-ViT-L/2)	3.40	30万	1024
ImageNet 256x256 (U-ViT-H/2)	2.29	50万	1024
ImageNet 512x512 (U-ViT-L/4)	4.67	50万	1024
ImageNet 512x512 (U-ViT-H/4)	4.05	50万	1024
MS-COCO (U-ViT-S/2)	5.95	100万	256
MS-COCO (U-ViT-S/2, Deep)	5.48	100万	256

训练与评估前的准备

自编码器

从该链接下载 stable-diffusion 目录（其中包含由 Stable Diffusion 转换而来的图像自编码器）。 将下载的目录放置到本代码库中的 assets/stable-diffusion 文件夹下。 这些自编码器用于潜在扩散模型中。

数据

• ImageNet 64x64：将标准 ImageNet 数据集（包含 train 和 val 目录）放入 assets/datasets/ImageNet。
• ImageNet 256x256 和 ImageNet 512x512：根据 scripts/extract_imagenet_feature.py 提取 ImageNet 特征。
• MS-COCO：下载 COCO 2014 训练、验证数据以及 标注。然后分别使用 scripts/extract_mscoco_feature.py、scripts/extract_test_prompt_feature.py 和 scripts/extract_empty_feature.py 提取相关特征。

FID 参考统计

从该链接下载 fid_stats 目录（其中包含 FID 的参考统计信息）。 将下载的目录放置到本代码库中的 assets/fid_stats 文件夹下。 除了用于评估外，这些参考统计信息还用于在训练过程中监控 FID。

训练

我们使用 huggingface accelerate 库来支持分布式数据并行和混合精度训练。以下是训练命令：

# 训练设置
num_processes=2  # 你拥有的 GPU 数量，例如 2
train_script=train.py  # 训练脚本，可选 <train.py|train_ldm.py|train_ldm_discrete.py|train_t2i_discrete.py>
                       # train.py：在像素空间进行训练
                       # train_ldm.py：在潜在空间使用连续时间步长进行训练
                       # train_ldm_discrete.py：在潜在空间使用离散时间步长进行训练
                       # train_t2i_discrete.py：在潜在空间进行文本到图像的训练
config=configs/cifar10_uvit_small.py  # 训练配置文件
                                      # 你可以通过修改配置文件来调整其他超参数

# 启动训练
accelerate launch --multi_gpu --num_processes $num_processes --mixed_precision fp16 $train_script --config=$config

我们提供了论文中复现 U-ViT 训练的所有命令：

# CIFAR10 (U-ViT-S/2)
accelerate launch --multi_gpu --num_processes 4 --mixed_precision fp16 train.py --config=configs/cifar10_uvit_small.py

# CelebA 64x64 (U-ViT-S/4)
accelerate launch --multi_gpu --num_processes 4 --mixed_precision fp16 train.py --config=configs/celeba64_uvit_small.py 

# ImageNet 64x64 (U-ViT-M/4)
accelerate launch --multi_gpu --num_processes 8 --mixed_precision fp16 train.py --config=configs/imagenet64_uvit_mid.py

# ImageNet 64x64 (U-ViT-L/4)
accelerate launch --multi_gpu --num_processes 8 --mixed_precision fp16 train.py --config=configs/imagenet64_uvit_large.py

# ImageNet 256x256 (U-ViT-L/2)
accelerate launch --multi_gpu --num_processes 8 --mixed_precision fp16 train_ldm.py --config=configs/imagenet256_uvit_large.py

# ImageNet 256x256 (U-ViT-H/2)
accelerate launch --multi_gpu --num_processes 8 --mixed_precision fp16 train_ldm_discrete.py --config=configs/imagenet256_uvit_huge.py

# ImageNet 512x512 (U-ViT-L/4)
accelerate launch --multi_gpu --num_processes 8 --mixed_precision fp16 train_ldm.py --config=configs/imagenet512_uvit_large.py

# ImageNet 512x512 (U-ViT-H/4)
accelerate launch --multi_gpu --num_processes 8 --mixed_precision fp16 train_ldm_discrete.py --config=configs/imagenet512_uvit_huge.py

# MS-COCO (U-ViT-S/2)
accelerate launch --multi_gpu --num_processes 4 --mixed_precision fp16 train_t2i_discrete.py --config=configs/mscoco_uvit_small.py

# MS-COCO (U-ViT-S/2, Deep)
accelerate launch --multi_gpu --num_processes 4 --mixed_precision fp16 train_t2i_discrete.py --config=configs/mscoco_uvit_small.py --config.nnet.depth=16

评估（计算 FID）

我们使用 huggingface accelerate 库来进行高效的混合精度和多 GPU 推理。以下是评估命令：

# 评估设置
num_processes=2  # 您拥有的 GPU 数量，例如 2
eval_script=eval.py  # 评估脚本，可选 <eval.py|eval_ldm.py|eval_ldm_discrete.py|eval_t2i_discrete.py>
                     # eval.py: 用于使用 train.py 训练的模型（即像素空间模型）
                     # eval_ldm.py: 用于使用 train_ldm.py 训练的模型（即具有连续时间步的潜在空间模型）
                     # eval_ldm_discrete.py: 用于使用 train_ldm_discrete.py 训练的模型（即具有离散时间步的潜在空间模型）
                     # eval_t2i_discrete.py: 用于使用 train_t2i_discrete.py 训练的模型（即潜在空间上的文本到图像模型）
config=configs/cifar10_uvit_small.py  # 训练配置

# 启动评估
accelerate launch --multi_gpu --num_processes $num_processes --mixed_precision fp16 eval_script --config=$config

生成的图片会存储在一个临时目录中，评估结束后会被删除。如果您想保留这些图片，请设置 --config.sample.path=/save/dir。

我们提供了复现论文中 FID 结果的所有命令：

# CIFAR10 (U-ViT-S/2)
accelerate launch --multi_gpu --num_processes 4 --mixed_precision fp16 eval.py --config=configs/cifar10_uvit_small.py --nnet_path=cifar10_uvit_small.pth

# CelebA 64x64 (U-ViT-S/4)
accelerate launch --multi_gpu --num_processes 4 --mixed_precision fp16 eval.py --config=configs/celeba64_uvit_small.py --nnet_path=celeba64_uvit_small.pth

# ImageNet 64x64 (U-ViT-M/4)
accelerate launch --multi_gpu --num_processes 8 --mixed_precision fp16 eval.py --config=configs/imagenet64_uvit_mid.py --nnet_path=imagenet64_uvit_mid.pth

# ImageNet 64x64 (U-ViT-L/4)
accelerate launch --multi_gpu --num_processes 8 --mixed_precision fp16 eval.py --config=configs/imagenet64_uvit_large.py --nnet_path=imagenet64_uvit_large.pth

# ImageNet 256x256 (U-ViT-L/2)
accelerate launch --multi_gpu --num_processes 8 --mixed_precision fp16 eval_ldm.py --config=configs/imagenet256_uvit_large.py --nnet_path=imagenet256_uvit_large.pth

# ImageNet 256x256 (U-ViT-H/2)
accelerate launch --multi_gpu --num_processes 8 --mixed_precision fp16 eval_ldm_discrete.py --config=configs/imagenet256_uvit_huge.py --nnet_path=imagenet256_uvit_huge.pth

# ImageNet 512x512 (U-ViT-L/4)
accelerate launch --multi_gpu --num_processes 8 --mixed_precision fp16 eval_ldm.py --config=configs/imagenet512_uvit_large.py --nnet_path=imagenet512_uvit_large.pth

# ImageNet 512x512 (U-ViT-H/4)
accelerate launch --multi_gpu --num_processes 8 --mixed_precision fp16 eval_ldm_discrete.py --config=configs/imagenet512_uvit_huge.py --nnet_path=imagenet512_uvit_huge.pth

# MS-COCO (U-ViT-S/2)
accelerate launch --multi_gpu --num_processes 4 --mixed_precision fp16 eval_t2i_discrete.py --config=configs/mscoco_uvit_small.py --nnet_path=mscoco_uvit_small.pth

# MS-COCO (U-ViT-S/2, Deep)
accelerate launch --multi_gpu --num_processes 4 --mixed_precision fp16 eval_t2i_discrete.py --config=configs/mscoco_uvit_small.py --config.nnet.depth=16 --nnet_path=mscoco_uvit_small_deep.pth

参考文献

如果您觉得该代码对您的研究有帮助，请考虑引用以下文献：

@inproceedings{bao2022all,
  title={All are Worth Words: A ViT Backbone for Diffusion Models},
  author={Bao, Fan and Nie, Shen and Xue, Kaiwen and Cao, Yue and Li, Chongxuan and Su, Hang and Zhu, Jun},
  booktitle = {CVPR},
  year={2023}
}

本实现基于以下项目：

• Extended Analytic-DPM（提供 CIFAR10 和 CelebA 64x64 的 FID 参考统计数据）
• guided-diffusion（提供 ImageNet 的 FID 参考统计数据）
• pytorch-fid（提供 FID 的 PyTorch 官方实现）
• dpm-solver（提供采样器）

U-ViT 快速上手指南

U-ViT 是一个基于 Vision Transformer (ViT) 的扩散模型骨干网络架构。它将时间、条件和噪声图像块均视为 Token，并采用长跳跃连接（Long Skip Connections），在图像生成任务中表现优异，尤其在 ImageNet 和 MS-COCO 数据集上取得了领先的 FID 分数。

1. 环境准备

• 操作系统: Linux (推荐)
• Python: 3.8+
• GPU: 支持 CUDA 的 NVIDIA GPU (训练高分辨率模型如 ImageNet 256x256+ 建议使用 A100 或多卡环境)
• PyTorch: 1.13.1+ (需匹配 CUDA 版本，官方示例基于 cu116)

2. 安装步骤

2.1 安装基础依赖

首先安装 PyTorch 及核心依赖库。国内用户可使用清华或阿里镜像加速下载。

# 安装 PyTorch (以 CUDA 11.6 为例，其他版本请参照 pytorch 官网)
pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu116

# 安装其他核心依赖
pip install accelerate==0.12.0 absl-py ml_collections einops wandb ftfy==6.1.1 transformers==4.23.1 timm==0.3.2

注意: timm 库可能需要应用一个补丁以兼容 PyTorch 1.8.1+，如遇报错请参考原仓库 issue #420。

2.2 安装加速组件（强烈推荐）

安装 xformers 可显著加速注意力计算并降低显存占用，是训练大模型的关键。

# 安装 xformers 和 triton
pip install -U xformers
pip install -U --pre triton

2.3 准备预训练资源与数据

在运行训练或评估前，需下载以下资源并放置于项目根目录的 assets/ 文件夹下：

1. 自编码器 (Autoencoder): 用于潜在空间扩散模型 (LDM)。
   • 下载 stable-diffusion 目录，重命名为 assets/stable-diffusion。
2. FID 统计参考值: 用于评估和训练监控。
   • 下载 fid_stats 目录，重命名为 assets/fid_stats。
3. 数据集:
   • ImageNet 64x64: 放入 assets/datasets/ImageNet。
   • ImageNet 256/512 & MS-COCO: 需先下载原始数据，然后运行 scripts/ 下的特征提取脚本预处理。

3. 基本使用

本项目使用 Hugging Face accelerate 库进行分布式训练和混合精度推理。

3.1 启动训练

以下命令展示了如何启动一个典型的训练任务（以 CIFAR10 为例）。你可以根据需求更换配置文件 (--config) 和训练脚本。

# 设置 GPU 数量 (例如 4 卡)
num_processes=4

# 启动训练 (CIFAR10, U-ViT-S/2)
# train.py 用于像素空间训练
# train_ldm_discrete.py 用于潜在空间离散时间步训练
# train_t2i_discrete.py 用于文生图训练
accelerate launch --multi_gpu --num_processes $num_processes --mixed_precision fp16 train.py --config=configs/cifar10_uvit_small.py

常用训练配置示例：

• ImageNet 256x256 (U-ViT-H/2):

  accelerate launch --multi_gpu --num_processes 8 --mixed_precision fp16 train_ldm_discrete.py --config=configs/imagenet256_uvit_huge.py
  

• MS-COCO 文生图:

  accelerate launch --multi_gpu --num_processes 4 --mixed_precision fp16 train_t2i_discrete.py --config=configs/mscoco_uvit_small.py
  

3.2 模型评估 (计算 FID)

使用提供的评估脚本计算生成图像的 FID 分数。

# 设置 GPU 数量
num_processes=2

# 评估像素空间模型
accelerate launch --multi_gpu --num_processes $num_processes --mixed_precision fp16 eval.py --config=configs/cifar10_uvit_small.py --ckpt_path=<路径到检查点文件>

# 评估潜在空间模型 (LDM)
accelerate launch --multi_gpu --num_processes $num_processes --mixed_precision fp16 eval_ldm_discrete.py --config=configs/imagenet256_uvit_huge.py --ckpt_path=<路径到检查点文件>

3.3 快速体验 (Colab)

如果你没有本地 GPU 环境，可以直接使用官方提供的 Colab Notebook 在云端采样 ImageNet 图像 (FID=2.29 模型)：