Transformer 和扩散模型的生成式 AI 实用指南（预览版）_扩散模型和transformer

原文：Hands-On Generative AI with Transformers and Diffusion Models

译者：飞龙

协议：CC BY-NC-SA 4.0

第一章：扩散模型

在 2020 年末，一个名为扩散模型的鲜为人知的模型类别开始在机器学习领域引起轰动。研究人员找出了如何使用这些模型生成比以前技术产生的合成图像质量更高的图像。随后出现了一系列论文，提出了改进和修改，进一步提高了质量。到 2021 年底，出现了像 GLIDE 这样的模型，展示了在文本到图像任务上令人难以置信的结果，几个月后，这些模型已经进入了主流，如 DALL-E 2 和 Stable Diffusion 等工具，使任何人都可以通过输入所需看到的文本描述来生成图像。

在本章中，我们将深入了解这些模型的工作原理。我们将概述使它们如此强大的关键见解，使用现有模型生成图像，以了解它们的工作方式，然后训练我们自己的模型，以进一步加深对此的理解。该领域仍在快速发展，但本章涵盖的主题应该为您打下坚实的基础。第五章将通过一个名为 Stable Diffusion 的模型，探索更高级的技术，第六章将探讨这些技术在简单图像生成之外的应用。

关键见解：迭代细化

那么，是什么使扩散模型如此强大呢？先前的技术，如 VAEs 或 GANs，通过模型的单次前向传递生成最终输出。这意味着模型必须在第一次尝试时就做到一切正确。如果它犯了一个错误，就无法返回并修复它。另一方面，扩散模型通过迭代多个步骤生成其输出。这种“迭代细化”允许模型纠正之前步骤中的错误，并逐渐改进输出。为了说明这一点，让我们看一个扩散模型的示例。

我们可以使用 Hugging Face 扩散器库加载预训练模型。该管道可用于直接创建图像，但这并不能显示出底层发生了什么：

# Load the pipeline
image_pipe = DDPMPipeline.from_pretrained("google/ddpm-celebahq-256")
image_pipe.to(device);

# Sample an image
image_pipe().images[0]
1
2
3
4
5
6

我们可以逐步重新创建采样过程，以更好地了解模型生成图像时发生了什么。我们使用随机噪声初始化我们的样本 x，然后通过模型运行 30 步。在右侧，您可以看到模型对特定步骤的最终图像预测-请注意，最初的预测并不特别好！我们不是直接跳到最终预测的图像，而是只在预测的方向上稍微修改 x（显示在左侧）。然后，我们再次将这个新的、稍微更好的 x 通过模型进行下一步的处理，希望能产生稍微改进的预测，这可以用来进一步更新 x，依此类推。经过足够的步骤，模型可以生成一些令人印象深刻的逼真图像。

# The random starting point for a batch of 4 images
x = torch.randn(4, 3, 256, 256).to(device)

# Set the number of timesteps lower
image_pipe.scheduler.set_timesteps(num_inference_steps=30)

# Loop through the sampling timesteps
for i, t in enumerate(image_pipe.scheduler.timesteps):

    # Get the prediction given the current sample x and the timestep t
    with torch.no_grad():
        noise_pred = image_pipe.unet(x, t)["sample"]

    # Calculate what the updated sample should look like with the scheduler
    scheduler_output = image_pipe.scheduler.step(noise_pred, t, x)

    # Update x
    x = scheduler_output.prev_sample

    # Occasionally display both x and the predicted denoised images
    if i % 10 == 0 or i == len(image_pipe.scheduler.timesteps) - 1:
        fig, axs = plt.subplots(1, 2, figsize=(12, 5))

        grid = torchvision.utils.make_grid(x, nrow=4).permute(1, 2, 0)
        axs[0].imshow(grid.cpu().clip(-1, 1) * 0.5 + 0.5)
        axs[0].set_title(f"Current x (step {i})")

        pred_x0 = scheduler_output.pred_original_sample
        grid = torchvision.utils.make_grid(pred_x0, nrow=4).permute(1, 2, 0)
        axs[1].imshow(grid.cpu().clip(-1, 1) * 0.5 + 0.5)
        axs[1].set_title(f"Predicted denoised images (step {i})")
        plt.show()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32

注意

如果那段代码看起来有点吓人-我们将在本章中解释这一切是如何工作的。现在，只需专注于结果。

学习如何逐渐改进“损坏”的输入的核心思想可以应用于广泛的任务范围。在本章中，我们将专注于无条件图像生成-也就是说，生成类似训练数据的图像，而不对这些生成的样本的外观进行额外的控制。扩散模型也已应用于音频、视频、文本等。虽然大多数实现使用我们将在这里介绍的某种“去噪”方法的变体，但正在出现利用不同类型的“损坏”以及迭代细化的新方法，这可能会使该领域超越目前专注于去噪扩散的焦点。令人兴奋的时刻！

训练扩散模型

在本节中，我们将从头开始训练一个扩散模型，以更好地了解它们的工作原理。我们将首先使用 Hugging Face diffusers 库中的组件。随着本章的进行，我们将逐渐揭开每个组件的工作原理。与其他类型的生成模型相比，训练扩散模型相对简单。我们反复进行：

• 从训练数据中加载一些图像。

• 以不同的量添加噪音。记住，我们希望模型能够很好地估计如何“修复”（去噪）极其嘈杂的图像和接近完美的图像。

• 将输入的嘈杂版本馈送到模型中。

• 评估模型在去噪这些输入时的表现。

• 使用这些信息来更新模型权重。

要使用经过训练的模型生成新图像，我们从完全随机的输入开始，并通过模型重复地将其馈送到模型中，在每次迭代中根据模型预测更新输入的小量。正如我们将看到的，有许多采样方法试图简化这个过程，以便我们可以尽可能少的步骤生成好的图像。

数据

对于这个例子，我们将使用来自 Hugging Face Hub 的图像数据集-具体来说，这个包含 1000 张蝴蝶图片的集合。在项目部分，您将看到如何使用自己的数据。

dataset = load_dataset("huggan/smithsonian_butterflies_subset", split="train")
1

在使用这些数据训练模型之前，我们需要做一些准备。图像通常表示为一个“像素”网格，每个像素有三个颜色通道（红色、绿色和蓝色）的颜色值在 0 到 255 之间。为了处理这些图像并使它们准备好进行训练，我们需要： - 将它们调整为固定大小 - （可选）通过随机水平翻转来添加一些增强，有效地使我们的数据集大小加倍 - 将它们转换为 PyTorch 张量（表示颜色值为 0 到 1 之间的浮点数） - 将它们标准化为具有均值为 0 的值，值在-1 到 1 之间

我们可以使用torchvision.transforms来完成所有这些操作：

image_size = 64

# Define data augmentations
preprocess = transforms.Compose(
    [
        transforms.Resize((image_size, image_size)),  # Resize
        transforms.RandomHorizontalFlip(),  # Randomly flip (data augmentation)
        transforms.ToTensor(),  # Convert to tensor (0, 1)
        transforms.Normalize([0.5], [0.5]),  # Map to (-1, 1)
    ]
)
1
2
3
4
5
6
7
8
9
10
11

接下来，我们需要创建一个数据加载器，以便加载应用了这些转换的数据批次：

batch_size = 32

def transform(examples):
    images = [preprocess(image.convert("RGB")) for image in examples["image"]]
    return {"images": images}

dataset.set_transform(transform)

train_dataloader = torch.utils.data.DataLoader(
    dataset, batch_size=batch_size, shuffle=True
)
1
2
3
4
5
6
7
8
9
10
11

我们可以通过加载单个批次并检查图像来检查这是否有效。

batch = next(iter(train_dataloader))
print('Shape:', batch['images'].shape,
      '\nBounds:', batch['images'].min().item(), 'to', batch['images'].max().item())
show_images(batch['images'][:8]*0.5 + 0.5) # NB: we map back to (0, 1) for display
1
2
3
4

Shape: torch.Size([32, 3, 64, 64])
Bounds: -0.9921568632125854 to 1.0
1
2

添加噪音

我们如何逐渐破坏我们的数据？最常见的方法是向图像添加噪音。我们添加的噪音量由噪音时间表控制。不同的论文和方法以不同的方式处理这个问题，我们将在本章后面进行探讨。现在，让我们看看一种常见的方法，基于 Ho 等人的论文“去噪扩散概率模型”。在扩散器库中，添加噪音是由称为调度器的东西处理的，它接收一批图像和一个“时间步长”列表，并确定如何创建这些图像的嘈杂版本：

scheduler = DDPMScheduler(num_train_timesteps=1000, beta_start=0.001, beta_end=0.02)
timesteps = torch.linspace(0, 999, 8).long()

x = batch['images'][:8]
noise = torch.rand_like(x)
noised_x = scheduler.add_noise(x, noise, timesteps)
show_images((noised_x*0.5 + 0.5).clip(0, 1))
1
2
3
4
5
6
7

在训练过程中，我们将随机选择时间步长。调度器接收一些参数（beta_start 和 beta_end），它用这些参数来确定给定时间步长应该存在多少噪音。我们将在第 X 节中更详细地介绍调度器。

UNet

UNet 是一种卷积神经网络，用于诸如图像分割之类的任务，其中期望的输出与输入具有相同的空间范围。它由一系列“下采样”层组成，用于减小输入的空间大小，然后是一系列“上采样”层，用于再次增加输入的空间范围。下采样层通常也后面跟着一个“跳跃连接”，将下采样层的输出连接到上采样层的输入。这允许上采样层“看到”网络早期的高分辨率表示，这对于具有图像样式输出的任务特别有用，其中这些高分辨率信息尤为重要。

扩散库中使用的 UNet 架构比 2015 年 Ronneberger 等人提出的原始 UNet更先进，增加了注意力和残差块等功能。我们稍后会更仔细地看一下，但这里的关键特点是它可以接受一个输入（嘈杂的图像）并产生一个形状相同的预测（预测的噪音）。对于扩散模型，UNet 通常还接受时间步作为额外的条件，我们将在 UNet 深入探讨部分再次探讨这一点。

这是我们如何创建一个 UNet 并将我们的一批嘈杂图像输入其中的方法：

# Create a UNet2DModel
model = UNet2DModel(
    in_channels=3,  # 3 channels for RGB images
    sample_size=64,  # Specify our input size
    block_out_channels=(64, 128, 256, 512), # N channels per layer
    down_block_types=("DownBlock2D", "DownBlock2D",
                      "AttnDownBlock2D", "AttnDownBlock2D"),
    up_block_types=("AttnUpBlock2D", "AttnUpBlock2D",
                    "UpBlock2D", "UpBlock2D"),
)

# Pass a batch of data through
with torch.no_grad():
    out = model(noised_x, timestep=timesteps).sample
out.shape
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

torch.Size([8, 3, 64, 64])
1

注意输出与输入的形状相同，这正是我们想要的。

训练

现在我们的模型和数据准备好了，我们可以开始训练。我们将使用学习率为 3e-4 的 AdamW 优化器。对于每个训练步骤，我们：

• 加载一批图像。

• 向图像添加噪音，选择随机时间步来确定添加多少噪音。

• 将嘈杂图像输入模型。

• 计算损失，即模型预测与目标之间的均方误差 - 在这种情况下是我们添加到图像中的噪音。这被称为噪音或“epsilon”目标。您可以在第 X 节中找到有关不同训练目标的更多信息。

• 反向传播损失并使用优化器更新模型权重。

在代码中，所有这些看起来是这样的：

num_epochs = 50 # How many runs through the data should we do?
lr = 1e-4 # What learning rate should we use
model = model.to(device) # The model we're training (defined in the previous section)
optimizer = torch.optim.AdamW(model.parameters(), lr=lr) # The optimizer
losses = [] # somewhere to store the loss values for later plotting

# Train the model (this takes a while!)
for epoch in range(num_epochs):
    for step, batch in enumerate(train_dataloader):

        # Load the input images
        clean_images = batch["images"].to(device)

        # Sample noise to add to the images
        noise = torch.randn(clean_images.shape).to(clean_images.device)

        # Sample a random timestep for each image
        timesteps = torch.randint(
            0,
            scheduler.num_train_timesteps,
            (clean_images.shape[0],),
            device=clean_images.device,
        ).long()

        # Add noise to the clean images according timestep
        noisy_images = scheduler.add_noise(clean_images, noise, timesteps)

        # Get the model prediction for the noise
        noise_pred = model(noisy_images, timesteps, return_dict=False)[0]

        # Compare the prediction with the actual noise:
        loss = F.mse_loss(noise_pred, noise)

        # Store the loss for later plotting
        losses.append(loss.item())

        # Update the model parameters with the optimizer based on this loss
        loss.backward(loss)
        optimizer.step()
        optimizer.zero_grad()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40

在 GPU 上运行上述代码大约需要一个小时，所以在等待时喝杯茶或者减少时代的数量。训练后损失曲线如下所示：

# Plot the loss curve:
plt.plot(losses);
1
2

随着模型学习去噪图像，损失曲线呈下降趋势。由于根据每次迭代中随机时间步的随机抽样向图像添加不同数量的噪音，曲线相当嘈杂。仅通过观察噪音预测的均方误差很难判断这个模型是否能够很好地生成样本，所以让我们继续下一节，看看它的表现如何。

采样

扩散库使用了“管道”的概念，将生成扩散模型样本所需的所有组件捆绑在一起：

pipeline = DDPMPipeline(unet=model, scheduler=scheduler)
ims = pipeline(batch_size=4).images
show_images(ims, nrows=1)
1
2
3

当然，将创建样本的工作交给管道并不能真正展示出我们正在进行的工作。因此，这里有一个简单的采样循环，展示了模型如何逐渐改进输入图像：

# Random starting point (4 random images):
sample = torch.randn(4, 3, 64, 64).to(device)

for i, t in enumerate(scheduler.timesteps):

    # Get model pred
    with torch.no_grad():
        noise_pred = model(sample, t).sample

    # Update sample with step
    sample = scheduler.step(noise_pred, t, sample).prev_sample

show_images(sample.clip(-1, 1)*0.5 + 0.5, nrows=1)
1
2
3
4
5
6
7
8
9
10
11
12
13

这与我们在本章开头使用的代码相同，用来说明迭代改进的概念，但希望现在你对这里发生的事情有了更好的理解。我们从一个完全随机的输入开始，然后在一系列步骤中由模型进行改进。每一步都是对输入的小更新，基于模型对该时间步的噪音的预测。我们仍然在pipeline.scheduler.step()的调用背后抽象了一些复杂性 - 在后面的章节中，我们将更深入地探讨不同的采样方法以及它们的工作原理。

评估

可以使用 FID 分数（Fréchet Inception Distance）来评估生成模型的性能。FID 分数通过比较从预训练的神经网络中提取的特征图之间的统计数据，衡量生成样本与真实样本的相似程度。分数越低，给定模型生成的图像的质量和逼真度就越好。FID 分数因其能够提供对不同类型生成网络的“客观”比较指标而受到欢迎，而无需依赖人类判断。

尽管 FID 分数很方便，但也有一些重要的注意事项需要注意：

• 给定模型的 FID 分数取决于用于计算它的样本数量，因此在模型之间进行比较时，我们需要确保报告的分数都是使用相同数量的样本计算的。通常做法是在这个目的上使用 50,000 个样本，尽管为了节省时间，您可能在开发过程中评估较少数量的样本，只有在准备发布结果时才进行完整评估。

• 在计算 FID 时，图像被调整为 299 像素的正方形图像。这使得它对于极低分辨率或高分辨率图像的度量变得不太有用。不同的深度学习框架处理调整大小的方式也有细微差异，这可能导致 FID 分数有小的差异！我们建议使用clean-fid这样的库来标准化 FID 计算。

• 用作 FID 特征提取器的网络通常是在 Imagenet 分类任务上训练的模型。在生成不同领域的图像时，这个模型学到的特征可能不太有用。更准确的方法是在特定领域的数据上训练一个分类网络，但这会使得在不同论文和方法之间比较分数变得更加困难，所以目前 Imagenet 模型是标准选择。

• 如果您保存生成的样本以供以后评估，格式和压缩也会影响 FID 分数。尽量避免低质量的 JPEG 图像。

即使考虑了所有这些警告，FID 分数只是质量的粗略度量，不能完美地捕捉使图像看起来更“真实”的微妙之处。因此，用它们来了解一个模型相对于另一个模型的表现，但也要看看每个模型生成的实际图像，以更好地了解它们的比较。人类偏好仍然是质量的黄金标准，最终这是一个相当主观的领域！

深入：噪音时间表

在上面的训练示例中，其中一步是“添加不同数量的噪音”。我们通过在 0 到 1000 之间选择一个随机时间步长来实现这一点，然后依靠调度程序添加适当数量的噪音。同样，在采样过程中，我们再次依靠调度程序告诉我们使用哪些时间步长以及如何根据模型预测从一个时间步长移动到下一个时间步长。事实证明，选择添加多少噪音是一个重要的设计决策，可以极大地影响给定模型的性能。在本节中，我们将看到为什么会这样，并探讨实践中使用的一些不同方法。

为什么要添加噪音？

在本章的开头，我们说扩散模型背后的关键思想是迭代改进。在训练期间，我们通过不同的方式“腐败”输入。在推断期间，我们从“最大程度上腐败”的输入开始，迭代地“去腐败”它，希望最终得到一个不错的最终结果。

到目前为止，我们已经专注于一种特定的“腐败”：添加高斯噪声。这样做的一个原因是扩散模型的理论基础 - 如果我们使用不同的腐败方法，我们在技术上就不再做“扩散”了！然而，Bansal 等人的一篇题为Cold Diffusion的论文戏剧性地表明，我们不一定需要仅仅出于理论上的便利而限制自己使用这种方法。他们表明，类似扩散模型的方法适用于许多不同的“腐败”方法（见图 1-1）。最近，像MUSE、MaskGIT和PAELLA这样的模型已经使用了随机标记屏蔽或替换作为等效的“腐败”方法，用于量化数据 - 也就是说，用离散标记而不是连续值表示的数据。

图 1-1。Cold Diffusion Paper 中使用的不同退化的示意图

尽管如此，出于几个原因，添加噪音仍然是最受欢迎的方法：

• 我们可以轻松控制添加的噪音量，从“完美”到“完全损坏”平稳过渡。这对于像减少图像分辨率这样的事情并不适用，这可能会导致“离散”的过渡。

• 我们可以有许多有效的随机起始点进行推断，不像一些方法可能只有有限数量的可能的初始（完全损坏）状态，比如完全黑色的图像或单像素图像。

所以，至少目前，我们将坚持添加噪音作为我们的损坏方法。接下来，让我们更仔细地看看如何向图像添加噪音。

开始简单

我们有一些图像（x），我们想以某种方式将它们与一些随机噪音结合起来。

x = next(iter(train_dataloader))['images'][:8]
noise = torch.rand_like(x)
1
2

我们可以线性插值（lerp）它们之间的一些量。这给我们一个函数，它在“量”从 0 到 1 变化时，从原始图像 x 平稳过渡到纯噪音：

def corrupt(x, noise, amount):
  amount = amount.view(-1, 1, 1, 1) # make sure it's broadcastable
  return x*(1-amount) + noise*amount
1
2
3

让我们在一批数据上看看这个过程，噪音的量从 0 到 1 变化：

amount = torch.linspace(0, 1, 8)
noised_x = corrupt(x, noise, amount)
show_images(noised_x*0.5 + 0.5)
1
2
3

这似乎正是我们想要的，从原始图像平稳过渡到纯噪音。现在，我们在这里创建了一个噪音时间表，它接受从 0 到 1 的“量”值。这被称为“连续时间”方法，我们在时间尺度上表示从 0 到 1 的完整路径。其他方法使用离散时间方法，使用一些大整数的“时间步长”来定义噪音调度器。我们可以将我们的函数封装成一个类，将连续时间转换为离散时间步长，并适当添加噪音：

class SimpleScheduler():
  def __init__(self):
    self.num_train_timesteps = 1000
  def add_noise(self, x, noise, timesteps):
    amount = timesteps / self.num_train_timesteps
    return corrupt(x, noise, amount)

scheduler = SimpleScheduler()
timesteps = torch.linspace(0, 999, 8).long()
noised_x = scheduler.add_noise(x, noise, timesteps)
show_images(noised_x*0.5 + 0.5)
1
2
3
4
5
6
7
8
9
10
11

现在我们有了一些可以直接与扩散库中使用的调度器进行比较的东西，比如我们在训练中使用的 DDPMScheduler。让我们看看它是如何比较的：

scheduler = DDPMScheduler(beta_end=0.01)
timesteps = torch.linspace(0, 999, 8).long()
noised_x = scheduler.add_noise(x, noise, timesteps)
show_images((noised_x*0.5 + 0.5).clip(0, 1))
1
2
3
4

数学

文献中有许多竞争的符号和方法。例如，一些论文将噪音时间表参数化为“连续时间”，其中 t 从 0（无噪音）到 1（完全损坏）- 就像我们在上一节中的corrupt函数一样。其他人使用“离散时间”方法，其中整数时间步长从 0 到某个大数 T，通常为 1000。可以像我们的SimpleScheduler类一样在这两种方法之间进行转换-只需确保在比较不同模型时保持一致。我们将在这里坚持使用离散时间方法。

深入研究数学的一个很好的起点是之前提到的 DDPM 论文。您可以在这里找到注释实现，这是一个很好的额外资源，可以帮助理解这种方法。

论文开始时指定了从时间步 t-1 到时间步 t 的单个噪音步骤。这是他们的写法：

q ( 声明：本文内容由网友自发贡献，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：【wpsshop博客】