访问官网
Github
文档
论文
仔细审视时间步长可使扩散模型训练速度提升三倍
如果你喜欢SpeeD，请在GitHub上给我们一个星标⭐以获取最新更新。
论文 | 项目主页 | Hugging Face
本仓库包含了题为"仔细审视时间步长可使扩散模型训练速度提升三倍"的研究论文的代码和实现细节。在这篇论文中，我们介绍了SpeeD，一种用于加速扩散模型训练的新方法。
作者
- 王凯2, 周钰坤1,2, 施明嘉2, 袁志航3, 尚宇章4, 彭晓江*1, 张汉旺5, 尤阳2
- 1深圳技术大学, 2新加坡国立大学, 3Infinigence-AI, 4伊利诺伊理工学院, 和 5南洋理工大学 [王凯、周钰坤和施明嘉对本工作贡献相同。我们将尽快更新此仓库。]
😮 亮点
我们的方法易于兼容，可以加速扩散模型的训练。
✒️ 动机
受以下对时间步长观察的启发，我们提出了重采样 + 重加权策略，如下所示。
仔细观察时间步长，我们发现时间步长可以分为三个区域：加速区、减速区和收敛区。收敛区对应时间步长的样本对训练的益处有限，而这些时间步长占用了最多的时间。从经验上看，这些样本的训练损失比其他两个区域的要低得多。
非对称采样：抑制收敛区时间步长的出现。
变化感知加权：在扩散过程中变化更快的时间步长被赋予更多权重。
在训练扩散模型时获取权重和采样t
def t_sample(self, n, device): if self.faster: t = torch.multinomial(self.p, n // 2 + 1, replacement=True).to(device) # 双重采样，可以平衡多步骤任务训练 dual_t = torch.where(t < self.meaningful_steps, self.meaningful_steps - t, t - self.meaningful_steps) t = torch.cat([t, dual_t], dim=0)[:n] weights = self.weights else: # 如果 t = torch.randint(0, self.num_timesteps, (n,), device=device) weights = None
return t, weights
你可以通过设置 diffusion.faster=True 来启用我们的加速模块。
# 配置文件
diffusion:
timestep_respacing: '250'
faster: true #启用训练加速模块
🛠️ 要求和安装
这个代码库不使用硬件加速技术，实验环境并不复杂。
你可以创建一个新的conda环境：
conda env create -f environment.yml
conda activate speed
或者通过以下方式安装必要的包：
pip install -r requirements.txt
如有必要，我们将提供更多方法（如docker）来方便配置实验环境。
🗝️ 教程
我们提供了一个完整的生成任务流程，包括训练、推理和测试。目前的代码仅兼容类别条件图像生成任务。我们将在未来兼容更多关于扩散的生成任务。
我们重构了facebookresearch/DiT的代码，并使用OmegaConf加载配置。配置文件加载规则是递归的，以便更容易修改参数。简单来说，后面路径中的文件将覆盖base.yaml中的先前设置。
你可以通过修改配置文件和命令行来修改实验设置。关于配置读取的更多细节写在configs/README.md中。
对于每个实验，你必须通过命令提供两个参数，
-c: 配置路径；
-p: 阶段，包括['train', 'inference', 'sample']。
训练 & 推理
基线
使用256x256 ImageNet数据集和DiT-XL/2模型的类别条件图像生成任务。
# 训练：训练扩散模型并保存检查点
torchrun --nproc_per_node=8 main.py -c configs/image/imagenet_256/base.yaml -p train
# 推理：生成用于测试的样本
torchrun --nproc_per_node=8 main.py -c configs/image/imagenet_256/base.yaml -p inference
# 采样：为可视化采样一些图像
python main.py -c configs/image/imagenet_256/base.yaml -p sample
消融
你可以通过修改配置文件和命令行来修改实验设置。关于配置的更多细节在configs/README.md中。
例如，通过命令行更改采样时的无分类器引导比例：
python main.py -c configs/image/imagenet_256/base.yaml -p sample guidance_scale=1.5
测试
测试生成任务需要推理的结果。关于测试的更多细节在evaluations中。
🔒 许可
本项目的大部分内容根据LICENSE文件中的Apache 2.0许可发布。
✏️引用
如果你发现我们的代码在你的研究中有用，请考虑给一个星星⭐和引用📝。
@article{wang2024closer,
title={A Closer Look at Time Steps is Worthy of Triple Speed-Up for Diffusion Model Training},
author={Kai Wang, Yukun Zhou, Mingjia Shi, Zhihang Yuan, Yuzhang Shang, Xiaojiang Peng, Hanwang Zhang and Yang You},
year={2024},
journal={arXiv preprint arXiv:2405.17403},
}
👍 致谢
我们感谢Tianyi Li、Yuchen Zhang、Yuxin Li、Zhaoyang Zeng和Yanqing Liu对这项工作的评论。Kai Wang（想法、写作、故事、演示）、Yukun Zhou（实现）和Mingjia Shi（理论、写作、演示）对这项工作做出了同等贡献。Xiaojiang Peng、Hanwang Zhang和Yang You是平等的指导。Xiaojiang Peng是通讯作者。
我们感谢以下杰出的工作和对开源的慷慨贡献。
