<img src="https://yellow-cdn.veclightyear.com/835a84d5/6930f4d4-477d-486f-abe1-4fa19c8f2ec3.png" width="150" style="margin-bottom: 0.2;"/> <h2 align="center">仔细审视时间步长可使扩散模型训练速度提升三倍 <h5 align="center"> 如果你喜欢SpeeD，请在GitHub上给我们一个星标⭐以获取最新更新。 </h2>

论文 | 项目主页 | 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是通讯作者。

我们感谢以下杰出的工作和对开源的慷慨贡献。