DiT原始项目

该项目仅针对DiT训练，并未包含VAE 的训练

项目地址

论文主页

Fast-DiT readme

该项目仅针对DiT训练，并未包含VAE 的训练

项目地址

该项目是基于论文Scalable Diffusion Models with Transformers的pytorch改进实现

包含：

• PyTorch的改进实现 和原始的DiT实现

• 在ImageNet (512x512 and 256x256)预训练的条件分类的 DiT 模型；

• 一个独立的 Hugging Face Space 和 Colab notebook，用于运行预训练的DiT-XL/2模型

• 改进的DiT 训练脚本 和一些 训练建议

启动
首先，下载开源代码

git clone https://github.com/chuanyangjin/fast-DiT.git
cd DiT

我们提供了 environment.yml 文件，可创建Conda 虚拟环境。

如果想要在本地CPU运行预训练的模型，可以在文件中删除cudatoolkit 和pytorch-cuda 相关的依赖项。

conda env create -f environment.yml
conda activate DiT

Sampling

预训练的DiT checkpoints

你可以使用预训练模型样例sample.py

预训练的DiT模型权重会根据使用的模型自动下载。

根据输入模型尺寸的不同，脚本中进行了不同的参数设置转换（256x256 and 512x512），比如针对512x512 DiT-XL/2 模型，你可以使用以下命令：

python sample.py --image-size 512 --seed 1

为了更加方便，我们的预训练模型也可直接进行下载：

DiT Model	Image Resolution	FID-50K	Inception Score	Gflops
XL/2	256x256	2.27	278.24	119
XL/2	512x512	3.04	240.82	525

自定义的DiT checkpoints

如果你想训练一个新的DiT 模型，可使用 train.py (see below)。

你可以增加一个参数–ckpt 使用你自己的checkpoint进行演示。比如运行一个 256x256 DiT-L/4 模型，可使用以下命令 :

If you’ve trained a new DiT model with train.py (see below), you can add the --ckpt argument to use your own checkpoint instead. For example, to sample from the EMA weights of a custom 256x256 DiT-L/4 model, run:

python sample.py --model DiT-L/4 --image-size 256 --ckpt /path/to/model.pt

Training

训练之前的准备

在一个GPU节点上抽取ImageNet 特征：

torchrun --nnodes=1 --nproc_per_node=1 extract_features.py --model DiT-XL/2 --data-path /path/to/imagenet/train --features-path /path/to/store/features

训练DiT

We provide a training script for DiT in train.py. This script can be used to train class-conditional DiT models, but it can be easily modified to support other types of conditioning.

我们提供了一个DiT的训练脚本 train.py。该训练脚本可以被用作训练条件分类的 DiT 模型，但是也可修改后用于其他类型的条件。

To launch DiT-XL/2 (256x256) training with 1 GPUs on one node:

accelerate launch --mixed_precision fp16 train.py --model DiT-XL/2 --features-path /path/to/store/features

To launch DiT-XL/2 (256x256) training with N GPUs on one node:

accelerate launch --multi_gpu --num_processes N --mixed_precision fp16 train.py --model DiT-XL/2 --features-path /path/to/store/features

或者，你可以选择提取并训练文件夹training options中的脚本。

PyTorch 训练结果

我们训练了 DiT-XL/2 和 DiT-B/4 模型

我们使用PyTorch训练脚本从头开始训练DiT-XL/2和DiT-B/4模型，以验证它重现了原始的JAX结果，达到数十万次训练迭代。

在我们的实验中，在合理的随机变化范围内，pytorch训练的模型与jax训练的模型相比，给出了类似(有时略好)的结果。一些数据点如下:

DiT Model	Train Steps	FID-50K (JAX Training)	FID-50K (PyTorch Training)	PyTorch Global Training Seed
XL/2	400K	19.5	18.1	42
B/4	400K	68.4	68.9	42
B/4	400K	68.4	68.3	100

这些模型在256x256分辨率下进行训练; 我们使用8x A100来训练XL/2，使用4x A100来训练B/4。注意，这里的FID是使用mse VAE解码器，在没有指导的情况下(cfg-scale=1)，通过250个DDPM采样步骤计算得到的。

改进训练效果

与原始实现相比，实现了一些训练加速和节省内存的特征，包括梯度检查点、混合精度训练和预提取的VAE特征，在DiT-XL/2上的速度提高了95%，内存减少了60%。一些数据点使用A100的全局批处理大小为128:

gradient checkpointing	mixed precision training	feature pre-extraction	training speed	memory
❌	❌	❌	-	out of memory
✔	❌	❌	0.43 steps/sec	44045 MB
✔	✔	❌	0.56 steps/sec	40461 MB
✔	✔	✔	0.84 steps/sec	27485 MB

Evaluation (FID, Inception Score, etc.)

我们提供一个’ sample_ddp.py '脚本，它可以并行地从DiT模型中对大量图像进行采样。

这个脚本生成一个样本文件夹和. npz文件一样，可以直接与ADM的TensorFlow评估套件一起使用，以计算FID, Inception分数和其他指标。

例如，要在N个gpu上从预训练的DiT-XL/2模型中采样50K张图像，运行:

torchrun --nnodes=1 --nproc_per_node=N sample_ddp.py --model DiT-XL/2 --num-fid-samples 50000

补充建议，详情请参见 sample_ddp.py

总结

开源项目仅能下载DiT-XL模型
微软开源了DiT-B模型，下载链接