DiTFastAttn 使用与启动教程

DiTFastAttn 使用与启动教程

DiTFastAttn 项目地址: https://gitcode.com/gh_mirrors/di/DiTFastAttn

1. 项目介绍

DiTFastAttn 是一种针对扩散变换器模型（Diffusion Transformers，DiT）的注意力压缩方法。DiT 在图像和视频生成方面表现出色，但由于自注意力机制的二次复杂度，其在计算上面临挑战。DiTFastAttn 通过识别和解决注意力计算中的三种冗余——空间冗余、时间冗余和条件冗余，来减轻 DiT 的计算瓶颈。项目提供了减少这些冗余的技巧，包括窗口注意力与残差缓存、时间相似性降低和条件冗余消除。

2. 项目快速启动

首先，确保安装以下环境：

conda create -n difa python=3.10

接着，安装所需的 Python 包：

pip install torch numpy packaging matplotlib scikit-image ninja
pip install git+https://github.com/huggingface/diffusers
pip install thop pytorch_fid torchmetrics accelerate torchmetrics[image] beautifulsoup4 ftfy flash-attn transformers SentencePiece

准备数据集：

• 对于 ImageNet 数据集，运行以下代码来采样真实图像：

python data/sample_real_images.py <imagenet_path>

• 如果使用 PixArt，请将 COCO 数据集放置在 data/mscoco 目录下。

启动 DiT 压缩：

python run_dit.py --n_calib 8 --n_steps 50 --window_size 128 --threshold 0.05 --eval_n_images 5000

启动 PixArt 1k 压缩：

python run_pixart.py --n_calib 6 --n_steps 50 --window_size 512 --threshold 0.0725 --eval_n_images 5000

启动 Opensora 压缩：

python run_opensora.py --threshold 0.05 --window_size 50 --n_calib 4 --use_cache

注意：在使用 Opensora 之前，需要根据 Opensora 的 README 进行安装。如果遇到问题，请切换到 Opensora 的 ea41df3d6cc5f38 提交。

3. 应用案例和最佳实践

• 案例一： 使用 DiTFastAttn 对 DiT 模型进行压缩，以减少计算资源和时间的需求。
• 案例二： 在图像生成任务中，通过调整 --window_size 和 --threshold 参数，找到最佳的性能与压缩比平衡点。
• 最佳实践： 在压缩前对模型进行充分校准（--n_calib 参数调整），以确保压缩后的模型仍能保持良好的生成质量。

4. 典型生态项目

DiTFastAttn 可以与以下开源项目结合使用，以构建更强大的图像和视频生成系统：

• Diffusers: 一个用于构建和训练扩散模型的库。
• Open-Sora: 一个用于加速神经网络的优化框架。

通过整合这些项目，可以进一步优化模型性能，并拓宽应用范围。

DiTFastAttn 项目地址: https://gitcode.com/gh_mirrors/di/DiTFastAttn