Skip to content

可复现的管道

扩散模型本质上是随机的,这使得它每次运行时都能生成不同的输出。但在某些时候,你希望每次都能生成相同的输出,比如在测试、复制结果,甚至提高图像质量时。虽然你不能期望在不同平台上获得完全相同的结果,但你可以在一定容差范围内期望在不同版本和平台之间获得可复现的结果(尽管这一点也可能有所不同)。

本指南将向你展示如何在 CPU 和 GPU 上控制随机性以实现确定性生成。

TIP

我们强烈建议阅读 PyTorch 关于可复现性的声明

"在 PyTorch 的不同版本、单个提交或不同平台上,完全可复现的结果无法保证。此外,即使使用相同的种子,CPU 和 GPU 执行之间的结果也可能无法复现。"

控制随机性

在推理过程中,管道高度依赖于随机采样操作,这些操作包括创建高斯噪声张量以去噪和在调度步骤中添加噪声。

查看在两次推理步骤后 [DDIMPipeline] 中的张量值。

python
from diffusers import DDIMPipeline
import numpy as np

ddim = DDIMPipeline.from_pretrained( "google/ddpm-cifar10-32", use_safetensors=True)
image = ddim(num_inference_steps=2, output_type="np").images
print(np.abs(image).sum())

运行上述代码会打印一个值,但如果你再次运行它,会得到一个不同的值。

每次运行管道时,torch.randn 都会使用不同的随机种子来创建高斯噪声张量。这导致每次运行时结果都不同,并使扩散管道能够每次生成不同的随机图像。

但如果你需要可靠地生成相同的图像,这取决于你是在 CPU 还是 GPU 上运行管道。

TIP

传递 Generator 对象而不是表示种子的整数值到管道中,这可能看起来不太直观。然而,这是在 PyTorch 中处理概率模型时推荐的设计,因为 Generator 是一个 随机状态,可以传递给多个管道。一旦 Generator 被消耗,状态 会就地改变,这意味着即使你将相同的 Generator 传递给不同的管道,也不会产生相同的结果,因为状态已经改变。

确定性算法

你还可以配置 PyTorch 使用确定性算法来创建可重现的管道。缺点是确定性算法可能比非确定性算法更慢,你可能会观察到性能下降。

非确定性行为发生在操作在多个 CUDA 流中启动时。为了避免这种情况,将环境变量 CUBLAS_WORKSPACE_CONFIG 设置为 :16:8,以在运行时仅使用一个缓冲区大小。

PyTorch 通常会基准测试多种算法以选择最快的算法,但如果你需要可重现性,应该禁用此功能,因为基准测试每次可能会选择不同的算法。设置 Diffusers enable_full_determinism 以启用确定性算法。

py
enable_full_determinism()

现在当你运行相同的管道两次时,你将获得相同的结果。

py
import torch
from diffusers import DDIMScheduler, StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", use_safetensors=True).to("cuda")
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
g = torch.Generator(device="cuda")

prompt = "A bear is playing a guitar on Times Square"

g.manual_seed(0)
result1 = pipe(prompt=prompt, num_inference_steps=50, generator=g, output_type="latent").images

g.manual_seed(0)
result2 = pipe(prompt=prompt, num_inference_steps=50, generator=g, output_type="latent").images

print("L_inf dist =", abs(result1 - result2).max())
"L_inf dist = tensor(0., device='cuda:0')"

确定性批量生成

创建可重复管道的一个实际应用是 确定性批量生成。你生成一批图像并选择一个图像,通过更详细的提示来改进它。主要思想是将一个 Generator 列表传递给管道,并将每个 Generator 绑定到一个种子,以便你可以重用它。

让我们使用 stable-diffusion-v1-5/stable-diffusion-v1-5 检查点并生成一批图像。

py
import torch
from diffusers import DiffusionPipeline
from diffusers.utils import make_image_grid

pipeline = DiffusionPipeline.from_pretrained(
    "stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
pipeline = pipeline.to("cuda")

定义四个不同的 Generator,并分别为每个 Generator 分配一个种子(03)。然后生成一批图像并选择其中一个进行迭代。

WARNING

使用列表推导式在 range() 中指定的批处理大小上进行迭代,为批处理中的每个图像创建一个唯一的 Generator 对象。如果你将 Generator 乘以批处理大小的整数,它只会创建 一个 Generator 对象,并依次用于批处理中的每个图像。

py
> [torch.Generator().manual_seed(seed)] * 4
> ```




```python
generator = [torch.Generator(device="cuda").manual_seed(i) for i in range(4)]
prompt = "Labrador in the style of Vermeer"
images = pipeline(prompt, generator=generator, num_images_per_prompt=4).images[0]
make_image_grid(images, rows=2, cols=2)

让我们改进第一张图片(你可以选择任何你想要的图片),该图片对应于 Generator 的种子 0。在提示中添加一些额外的文本,然后确保你使用相同的 Generator 和种子 0。所有生成的图片都应该与第一张图片相似。

python
prompt = [prompt + t for t in [", highly realistic", ", artsy", ", trending", ", colorful"]]
generator = [torch.Generator(device="cuda").manual_seed(0) for i in range(4)]
images = pipeline(prompt, generator=generator).images
make_image_grid(images, rows=2, cols=2)