ddpo-pytorch

ddpo-pytorch

这是一个在 PyTorch 中实现的 降噪扩散策略优化 (DDPO),支持 低秩适应 (LoRA)。与我们原始的研究代码 (您可以在 这里 找到) 不同,本实现在 GPU 上运行,如果启用 LoRA,则只需不到 10GB 的 GPU 内存就可以微调 Stable Diffusion!

安装

需要 Python 3.10 或更高版本。

git clone [email protected]:kvablack/ddpo-pytorch.git
cd ddpo-pytorch
pip install -e .

使用方法

accelerate launch scripts/train.py

这将立即开始使用 config/base.py 中的配置，在所有可用的 GPU 上对 Stable Diffusion v1.5 进行可压缩性的微调。只要每个 GPU 至少有 10GB 内存，就应该能正常运行。如果您不想登录 wandb，可以在上述命令之前运行 wandb disabled。

所有超参数的详细说明可以在 config/base.py 中找到。以下是一些重要的超参数。

从高层次来看,微调扩散模型的问题由 2 个因素定义:用于生成图像的提示集合以及用于评估这些图像的奖励函数。提示由 prompt_fn 定义,该函数不接受参数,每次调用时都会生成一个随机提示。奖励函数由 reward_fn 定义,该函数接受一批图像,并返回对这些图像的奖励。目前实现的所有提示和奖励函数均可在 ddpo_pytorch/prompts.py 和 ddpo_pytorch/rewards.py 中找到。

批量大小和累积步骤 每个 DDPO 时期包括生成一批图像,计算它们的奖励,然后对这些图像执行一些训练步骤。