Fast Diffusion Models with Transformers (fast-DiT) 使用教程

Fast Diffusion Models with Transformers (fast-DiT) 使用教程

fast-DiT Fast Diffusion Models with Transformers 项目地址: https://gitcode.com/gh_mirrors/fa/fast-DiT

1. 项目介绍

Fast Diffusion Models with Transformers (fast-DiT) 是一个基于 PyTorch 的开源项目，旨在提供高效且可扩展的扩散模型实现。该项目改进了原始的 DiT（Diffusion Models with Transformers）模型，并提供了预训练模型和训练脚本，支持在 ImageNet 数据集上进行图像生成任务。

主要特性：

• 改进的 PyTorch 实现：提供更高效的训练和采样流程。
• 预训练模型：包含在 ImageNet 上训练的 256x256 和 512x512 分辨率的 DiT 模型。
• 自包含的运行环境：提供 Hugging Face Space 和 Colab 笔记本，方便快速运行预训练模型。
• 高效的训练脚本：支持梯度检查点、混合精度训练和预提取 VAE 特征，显著提升训练速度和降低内存消耗。

2. 项目快速启动

环境配置

首先，克隆项目仓库并创建 Conda 环境：

git clone https://github.com/chuanyangjin/fast-DiT.git
cd fast-DiT
conda env create -f environment.yml
conda activate DiT

如果仅需要在 CPU 上运行预训练模型，可以移除 cudatoolkit 和 pytorch-cuda 依赖。

运行预训练模型

使用 sample.py 脚本从预训练模型中采样图像：

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

此命令将从 512x512 分辨率的 DiT-XL/2 模型中生成图像。其他参数如 --model 和 --ckpt 可用于指定不同模型或自定义检查点。

训练自定义模型

提取 ImageNet 特征

在单节点单 GPU 上提取特征：

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

启动训练

在单节点单 GPU 上训练 DiT-XL/2 (256x256) 模型：

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

在单节点 N 个 GPU 上训练：

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

3. 应用案例和最佳实践

应用案例

• 图像生成：利用预训练模型生成高质量的图像，适用于艺术创作、游戏开发等领域。
• 图像编辑：通过微调模型，实现图像修复、风格迁移等任务。
• 数据增强：生成合成数据，用于训练其他机器学习模型。

最佳实践

• 使用预训练模型：直接使用提供的预训练模型可以快速获得高质量结果。
• 混合精度训练：启用混合精度训练可以显著提升训练速度并减少内存消耗。
• 梯度检查点：在内存有限的情况下，使用梯度检查点技术可以有效降低内存占用。
• 预提取 VAE 特征：预先提取并存储 VAE 特征，可以加速训练过程。

4. 典型生态项目

• Hugging Face Transformers：提供丰富的 Transformer 模型库，可与 fast-DiT 结合使用。
• PyTorch Lightning：简化 PyTorch 训练流程，提升开发效率。
• TensorFlow Probability：提供概率模型和推理工具，可用于扩展 fast-DiT 的应用范围。

通过以上教程，您可以快速上手 fast-DiT 项目，并进行图像生成和模型训练等任务。更多详细信息和高级用法，请参考项目官方文档和代码仓库。

fast-DiT Fast Diffusion Models with Transformers 项目地址: https://gitcode.com/gh_mirrors/fa/fast-DiT