U-DiT项目使用与启动教程

U-DiT项目使用与启动教程

U-DiT [NeurIPS 2024] The official code of "U-DiTs: Downsample Tokens in U-Shaped Diffusion Transformers" 项目地址: https://gitcode.com/gh_mirrors/ud/U-DiT

1. 项目介绍

U-DiT（Downsample Tokens in U-Shaped Diffusion Transformers）是一个基于U-Net架构和Transformer的图像生成模型。该项目是NeurIPS 2024的官方代码实现，其主要创新点在于通过在Transformer中引入U-Net的降采样特性，以减少计算成本并提高模型性能。U-DiT-B模型能够在保持与DiT-XL/2相当性能的同时，仅使用其1/6的计算量（FLOPs）。

2. 项目快速启动

环境准备

在开始之前，您需要确保已经安装了以下依赖项：

pip install -r requirements.txt

数据准备

U-DiT模型可以使用原始的ImageNet数据集进行训练，或者使用VAE编码后的特征进行训练（推荐）。以下为数据准备的步骤：

• 使用原始ImageNet数据集：
  • 首先下载ImageNet数据集。
  • 使用以下命令进行训练：

torchrun --nnodes=1 --nproc_per_node=8 train.py --data-path={path to imagenet/train} --image-size=256 --model={model name} --epochs={iteration//5000}

• 使用VAE特征进行训练（推荐）：
  • 下载ImageNet的VAE特征。
  • 解压下载的文件。
  • 在训练命令中添加--feature-path={path to imagenet_feature}参数。

模型训练

• FP32训练：

torchrun --nnodes=1 --nproc_per_node=8 train.py --data-path {path to imagenet/train} --image-size=256 --model={model name} --epochs={iteration//5000}

• FP16训练：

accelerate launch --mixed_precision fp16 train_accelerate.py --data-path {path to imagenet/train} --image-size=256 --model={model name} --epochs={iteration//5000}

模型采样

使用以下命令进行并行采样：

torch --nnodes=1 --nproc_per_node=8 sample_ddp.py --ckpt={path to checkpoint} --image-size=256 --model={model name} --cfg-scale={cfg scale}

性能评估

从以下位置借用FID评估代码，并按照以下命令进行性能评估：

python evaluator.py {path to reference batch} {path to generated .npz}

3. 应用案例和最佳实践

U-DiT模型可以用于多种图像生成任务，例如图像修复、风格迁移等。最佳实践建议在训练前使用VAE特征以加快训练速度，并在模型训练时使用适当的降采样策略以提高效率。

4. 典型生态项目

U-DiT模型建立在DiT（Dual-Attention Transformer）的基础上，并与之兼容。它可以使用DiT的生态项目，如训练工具、模型权重和预训练特征等。此外，U-DiT社区也提供了一系列的Colab演示和长期训练结果的展示。

U-DiT [NeurIPS 2024] The official code of "U-DiTs: Downsample Tokens in U-Shaped Diffusion Transformers" 项目地址: https://gitcode.com/gh_mirrors/ud/U-DiT