可复现的管道
扩散模型本质上是随机的,这使得它每次运行都能生成不同的输出。但在某些时候,您希望每次都生成相同的输出,例如在测试、复制结果甚至提高图像质量时。虽然您不能期望在不同平台上获得完全相同的结果,但您可以在一定容差范围内期望在不同版本和平台之间获得可复现的结果(尽管即使这样也可能会有所不同)。
本指南将向您展示如何在 CPU 和 GPU 上控制随机性以进行确定性生成。
我们强烈建议阅读 PyTorch 的关于可复现性的声明
“PyTorch 版本、单个提交或不同平台之间不保证完全可复现的结果。此外,即使使用相同的种子,CPU 和 GPU 执行之间也可能无法复现结果。”
控制随机性
在推理过程中,管道严重依赖于随机采样操作,其中包括创建用于去噪的高斯噪声张量以及向调度步骤添加噪声。
查看DDIMPipeline在两个推理步骤后的张量值。
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 上运行管道。
将Generator
对象传递给管道而不是表示种子的整数可能看起来不直观。但是,在使用 PyTorch 中的概率模型时,这是推荐的设计,因为Generator
是一个可以依次传递给多个管道的随机状态。一旦Generator
被使用,状态就会就地更改,这意味着即使您将相同的Generator
传递给不同的管道,它也不会产生相同的结果,因为状态已经更改。
要在 CPU 上生成可复现的结果,您需要使用 PyTorch Generator 并设置种子。现在,当您运行代码时,它始终打印值为1491.1711
,因为带有种子的Generator
对象被传递给管道中的所有随机函数。您应该在您使用的任何硬件和 PyTorch 版本上获得类似的结果,即使不完全相同。
import torch
import numpy as np
from diffusers import DDIMPipeline
ddim = DDIMPipeline.from_pretrained("google/ddpm-cifar10-32", use_safetensors=True)
generator = torch.Generator(device="cpu").manual_seed(0)
image = ddim(num_inference_steps=2, output_type="np", generator=generator).images
print(np.abs(image).sum())
确定性算法
您还可以配置 PyTorch 使用确定性算法来创建可复现的管道。缺点是确定性算法可能比非确定性算法慢,您可能会观察到性能下降。
当在多个 CUDA 流中启动操作时,会出现非确定性行为。为避免这种情况,请将环境变量CUBLAS_WORKSPACE_CONFIG设置为:16:8
,以便在运行时仅使用一个缓冲区大小。
PyTorch 通常会对多种算法进行基准测试以选择最快的算法,但如果您希望获得可复现的结果,则应禁用此功能,因为基准测试可能会每次选择不同的算法。设置 Diffusers enable_full_determinism 以启用确定性算法。
enable_full_determinism()
现在,当您两次运行相同的管道时,您将获得相同的结果。
import torch
from diffusers import DDIMScheduler, StableDiffusionPipeline
pipe = StableDiffusionPipeline.from_pretrained("runwayml/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’s传递给管道并将每个Generator
绑定到一个种子,以便您可以重用它。
让我们使用runwayml/stable-diffusion-v1-5检查点并生成一批图像。
import torch
from diffusers import DiffusionPipeline
from diffusers.utils import make_image_grid
pipeline = DiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
pipeline = pipeline.to("cuda")
定义四个不同的Generator
并为每个Generator
分配一个种子(0
到3
)。然后生成一批图像并选择一个进行迭代。
使用列表推导式遍历range()
中指定的批次大小,为批次中的每个图像创建一个唯一的Generator
对象。如果您将Generator
乘以批次大小整数,它只会创建一个一个Generator
对象,该对象按顺序用于批次中的每个图像。
[torch.Generator().manual_seed(seed)] * 4
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)
让我们改进第一张图像(您可以选择任何您想要的图像),它对应于种子为0
的Generator
。在您的提示中添加一些其他文本,然后确保您重用种子为0
的相同Generator
。所有生成的图像都应该类似于第一张图像。
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)