SiT

使用可扩展插值变换器(SiT)探索基于流和扩散的生成模型
_{PyTorch官方实现}

论文 | 项目页面 |

本仓库包含我们论文中探索具有可扩展变换器的插值模型(SiTs)的PyTorch模型定义、预训练权重以及训练/采样代码。

使用可扩展插值变换器探索基于流和扩散的生成模型
马南野、Mark Goldstein、Michael Albergo、Nicholas Boffi、Eric Vanden-Eijnden、谢赛宁
纽约大学

我们提出了可扩展插值变换器(SiT)，这是一系列基于扩散变换器(DiT)骨干网络构建的生成模型。插值框架允许以比标准扩散模型更灵活的方式连接两个分布，使得可以对影响基于动态传输的生成模型的各种设计选择进行模块化研究：使用离散或连续时间学习、决定要学习的模型、选择连接分布的插值器以及部署确定性或随机采样器。通过仔细引入上述要素，SiT在条件ImageNet 256x256基准测试中使用完全相同的骨干网络、参数数量和GFLOPs，在所有模型规模上均超过了DiT。通过探索可以与学习分开调整的各种扩散系数，SiT实现了2.06的FID-50K分数。

本仓库包含：

• 🪐 SiT的简单PyTorch实现
• ⚡️ 在ImageNet 256x256上训练的预训练类别条件SiT模型
• 🛸 使用PyTorch DDP的SiT训练脚本

设置

首先，下载并设置仓库：

git clone https://github.com/willisma/SiT.git
cd SiT

我们提供了一个environment.yml文件，可用于创建Conda环境。如果你只想在本地CPU上运行预训练模型，可以从文件中删除cudatoolkit和pytorch-cuda要求。

conda env create -f environment.yml
conda activate SiT

采样

预训练SiT检查点。 你可以使用sample.py从我们的预训练SiT模型中进行采样。根据你使用的模型，预训练SiT模型的权重将自动下载。该脚本有各种参数可以调整采样器配置(ODE和SDE)、采样步骤、更改无分类器引导尺度等。例如，要使用默认ODE设置从我们的256x256 SiT-XL模型中进行采样，你可以使用：

python sample.py ODE --image-size 256 --seed 1

为方便起见，我们的预训练SiT模型也可以在这里直接下载：

自定义SiT检查点。 如果你使用train.py训练了新的SiT模型(参见下文)，你可以添加--ckpt参数来使用你自己的检查点。例如，要使用ODE采样器从自定义256x256 SiT-L/4模型的EMA权重中进行采样，请运行：

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

高级采样器设置

还有一些其他选项；详情请参阅train_utils.py。

训练SiT

我们在train.py中提供了SiT的训练脚本。要在一个节点上使用N个GPU启动SiT-XL/2（256x256）的训练：

torchrun --nnodes=1 --nproc_per_node=N train.py --model SiT-XL/2 --data-path /path/to/imagenet/train

日志记录。 要启用wandb，首先将WANDB_KEY、ENTITY和PROJECT设置为环境变量：

export WANDB_KEY="key"
export ENTITY="entity name"
export PROJECT="project name"

然后在训练命令中添加--wandb标志：

torchrun --nnodes=1 --nproc_per_node=N train.py --model SiT-XL/2 --data-path /path/to/imagenet/train --wandb

插值设置。 我们还支持不同的插值选择和模型预测。例如，要使用Linear插值和noise预测启动SiT-XL/2（256x256）：

torchrun --nnodes=1 --nproc_per_node=N train.py --model SiT-XL/2 --data-path /path/to/imagenet/train --path-type Linear --prediction noise

恢复训练。 从自定义检查点恢复训练：

torchrun --nnodes=1 --nproc_per_node=N train.py --model SiT-L/2 --data-path /path/to/imagenet/train --ckpt /path/to/model.pt

注意。 恢复训练将自动将模型、EMA和优化器状态以及训练配置恢复为与检查点相同的状态。

评估（FID、Inception Score等）

我们包含了一个sample_ddp.py脚本，可以并行从SiT模型中采样大量图像。该脚本生成一个包含样本的文件夹以及一个.npz文件，可以直接用于ADM的TensorFlow评估套件来计算FID、Inception Score和其他指标。例如，要在默认ODE采样器设置下使用N个GPU从我们预训练的SiT-XL/2模型中采样50K张图像，运行：

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

似然。 支持似然评估。要计算似然，您可以向ODE采样器添加--likelihood标志：

torchrun --nnodes=1 --nproc_per_node=N sample_ddp.py ODE --model SiT-XL/2 --likelihood

请注意，只有在ODE采样器下才能计算似然；有关更多详细信息和设置，请参见sample_ddp.py。

改进

通过以下方式可能会显著加快训练（和采样）速度：

• 在SiT模型中使用Flash Attention
• 使用PyTorch 2.0中的torch.compile

需要添加的基本功能：

• 监控FID和其他指标
• AMP/bfloat16支持

通过以下方式可能会提高似然计算的精度：

• 均匀/高斯去量化

与JAX的差异

我们的模型最初是在TPU上使用JAX训练的。此存储库中的权重直接从JAX模型移植而来。由于在不同平台（TPU与GPU）上采样，结果可能存在微小差异。我们观察到在TPU上采样的性能略差于GPU（FID为2.15，而论文中为2.06）。

许可证

该项目采用MIT许可证。详情请参见LICENSE。