本章内容包括
- 正向扩散和逆向扩散的工作原理
- 如何构建和训练去噪U-Net模型
- 使用训练好的U-Net生成花卉图像
- 文本生成图像Transformer的相关概念
- 使用DALL-E 2通过文本生成图像的Python程序编写
近年来,多模态大型语言模型(LLM)因其能够处理多种内容格式(如文本、图像、视频、音频和代码)而备受关注。文本生成图像Transformer是其中的典型代表,例如OpenAI的DALL-E 2、Google的Imagen和Stability AI的Stable Diffusion。这些模型可以根据文本描述生成高质量图像。
这些文本生成图像模型包含三个核心部分:文本编码器(将文本压缩成潜在表示)、将文本信息融入图像生成过程的方法,以及扩散机制(逐步优化图像,生成逼真输出)。理解扩散机制对掌握文本生成图像Transformer尤为重要,因为扩散模型构成了所有领先文本生成图像Transformer的基础。因此,本章首先从构建和训练一个扩散模型来生成花卉图像开始。你将深入理解正向扩散过程,即图像中逐步加入噪声直到变为随机噪声的过程。接着,你将训练模型学习逆转该过程,即从纯噪声图像开始,逐步去除噪声,直到生成类似训练集中的干净新图像。
扩散模型已成为生成高分辨率图像的首选技术。它们的成功在于能够模拟并逆转复杂的噪声添加过程,类似于对图像结构的深刻理解以及从抽象模式构建图像的能力。这种方法不仅保证了图像的高质量,还在多样性与准确性之间取得了平衡。
随后,我们将解释文本生成图像Transformer的工作原理。重点介绍OpenAI开发的对比语言-图像预训练(CLIP)模型,该模型能够理解并关联视觉与文本信息。CLIP处理两种输入:图像和文本(通常为标题或描述),分别通过两个编码器独立处理。
CLIP的图像分支使用视觉Transformer(ViT)将图像编码成高维向量空间,提取视觉特征;文本分支采用基于Transformer的语言模型,将文本描述编码到同一向量空间,捕捉语义特征。CLIP通过大量配对的图像与文本描述训练,使匹配对在向量空间中紧密对齐。
OpenAI的文本生成图像Transformer(如DALL-E 2)将CLIP作为核心组件。本章将教你如何获取OpenAI API密钥,并编写Python程序,使用DALL-E 2根据文本描述生成图像。
15.1 去噪扩散模型简介
扩散模型的概念可通过以下示例说明。设想利用扩散模型生成高分辨率花卉图像。首先收集一组高质量花卉图片作为训练数据。模型被指示逐步向这些图像添加少量随机噪声,这一过程称为正向扩散。经过多次噪声叠加后,训练图像最终变成纯随机噪声。接下来训练模型逆转该过程,从纯噪声图像出发,逐步去除噪声,直到生成与原始训练集图像无异的清晰图像。
训练完成后,模型接收随机噪声图像作为输入,经过多次迭代系统性地去除噪声,最终生成一张与训练集相似的高分辨率花卉图像。这就是扩散模型的基本原理。
本节将先探讨扩散模型的数学基础,然后深入U-Net架构——用于图像去噪与生成高分辨率花卉图像的模型。具体而言,U-Net采用了类似于第9至12章Transformer模型中见到的缩放点积注意力(SDPA)机制。最后,你将学习扩散模型的训练过程及训练后模型的图像生成流程。
15.1.1 正向扩散过程
已有多篇论文提出了具有相似机制的扩散模型。以花卉图像为例,图15.1展示了正向扩散过程的工作原理。
图15.1 正向扩散过程示意图。我们从训练集中一张干净的图像 x₀ 开始,向其添加噪声 ε₀,生成带噪图像 x₁ = √(1 – β₁) x₀ + √(β₁) ε₀。我们重复此过程共 1000 个时间步,直到图像 x₁₀₀₀ 变成纯随机噪声。
假设花卉图像 x₀(如图15.1左侧所示)服从分布 q(x)。在正向扩散过程中,我们将在 T = 1000 个步骤中逐步向图像添加少量噪声。噪声张量服从正态分布,且形状与花卉图像相同,为 (3, 64, 64),表示有三个颜色通道,图像的高和宽均为 64 像素。
扩散模型中的时间步长 扩散模型中的时间步长指的是在逐步向数据添加噪声并随后逆转该过程生成样本时的离散阶段。扩散模型的正向阶段会在一系列时间步长中逐渐添加噪声,将数据从其原始的干净状态转变为带噪分布。逆向阶段则在类似的时间步长序列中逆序进行,系统性地去除数据中的噪声,重建原始数据或生成新的高质量样本。逆向过程中的每个时间步都涉及预测对应正向步骤中添加的噪声,并将其减去,从而逐步去噪,直到恢复到干净状态。
在时间步 1 中,我们向图像 x₀ 添加噪声 ε₀,从而得到带噪图像 x₁:
也就是说,x₁ 是 x₀ 和 ε₀ 的加权和,其中 β₁ 表示噪声的权重。β 的值会随着时间步的不同而变化,因此有了下标 1。如果我们假设 x₀ 和 ε₀ 彼此独立,且都服从标准正态分布(即均值为 0,方差为 1),那么带噪图像 x₁ 也将服从标准正态分布。这一点很容易证明,因为……
和
我们可以在接下来的 T-1 个时间步继续向图像中添加噪声,使得……
我们可以使用重参数化技巧,并定义 以及…
允许我们在任意时间步 t 采样 ,其中 t 可以取值范围为 。那么我们有:
є 是 є₀、є₁、…、є_{t-1} 的组合,利用两个正态分布相加仍然是正态分布的性质得到的。有关证明,可以参考 Lilian Weng 在 mng.bz/Aalg 的博客。
图 15.1 左侧展示了训练集中一张干净的花卉图像 。在第一个时间步,我们向其注入噪声 є₀,形成带噪声的图像 (图中间图像)。我们重复此过程共 1000 个时间步,直到图像变成随机噪声(最右侧图像)。
15.1.2 使用 U-Net 模型去噪图像
理解了正向扩散过程后,我们来讨论逆向扩散过程(即去噪过程)。如果我们能训练一个模型逆转正向扩散过程,就可以给模型输入随机噪声,让它生成一张带噪声的花卉图像。然后再将该带噪声图像输入训练好的模型,生成更清晰(但仍带有噪声)的图像。我们可以重复这个过程多次,直到生成的图像干净且与训练集中的图像无法区分。逆向扩散过程采用多次推理步骤而非一步完成,是逐步从噪声分布重构高质量数据的关键。这样能更好地控制,保证生成的质量稳定且高质量。
为此,我们将构建一个去噪 U-Net 模型。U-Net 架构最初设计用于生物医学图像分割,特点是对称结构,由收缩路径(编码器)和扩展路径(解码器)组成,中间通过瓶颈层相连。对于去噪任务,U-Net 模型被调整用于在去除图像噪声的同时保留重要细节。U-Net 在去噪任务中优于简单卷积网络,因为它能高效捕捉图像的局部和全局特征。
图 15.2 展示了本章所用去噪 U-Net 的结构示意图。
模型输入带噪图像和对应的时间步(即公式 15.3 中的 和 t),输出预测该图像中的噪声(即 є)。由于带噪图像是原始干净图像与噪声的加权和(见公式 15.3),通过预测噪声,我们可以推断并重建原始图像。
收缩路径(编码器,图 15.2 左侧)由多个卷积层和池化层组成,逐步对图像进行下采样,提取并编码不同层次的特征。网络学会识别对去噪有用的模式和特征。
瓶颈层(图 15.2 底部)连接编码器和解码器,由卷积层组成,负责捕捉图像最抽象的表示。
扩展路径(解码器,图 15.2 右侧)由上采样层和卷积层组成,逐步对特征图进行上采样,重建图像,同时通过跳跃连接融合编码器提取的特征。跳跃连接(图中虚线所示)是 U-Net 模型的关键,它使模型能够通过结合低层和高层特征,保留输入图像的细节。接下来,我将简要说明跳跃连接的工作原理。
图 15.2 去噪 U-Net 模型的架构示意。U-Net 架构的特点是对称结构,由收缩路径(编码器)和扩展路径(解码器)组成,中间通过瓶颈层相连。该模型旨在去除图像噪声的同时,保留重要细节。模型的输入是一张带噪声的图像及其对应的时间步,输出是图像中预测的噪声。
在 U-Net 模型中,跳跃连接通过将编码器路径中的特征图与解码器路径中对应的特征图进行拼接实现。这些特征图通常具有相同的空间尺寸,但由于它们经历了不同路径的处理,可能有不同的表现。编码过程中,输入图像会逐步下采样,一些空间信息(如边缘和纹理)可能丢失。跳跃连接通过直接传递编码器的特征图到解码器,绕过信息瓶颈,帮助保留这些信息。
例如,图 15.2 顶部的虚线表示模型将编码器中 Conv2D 层输出的形状为 (128, 64, 64) 的特征图与解码器中 Conv2D 层输入的同样形状的特征图拼接,最终解码器 Conv2D 层的输入形状变为 (256, 64, 64)。
通过将解码器中的高级抽象特征与编码器中的低级细节特征相结合,跳跃连接使模型能够更好地重建去噪图像中的细节。这在去噪任务中尤为重要,因为保留图像中的细微细节至关重要。
缩放点积注意力(Scaled Dot Product Attention, SDPA)机制同时被应用在去噪 U-Net 的收缩路径和扩展路径的最后一个模块中,配合层归一化和残差连接(图 15.2 中标注为 Attn/Norm/Add)。该 SDPA 机制与第 9 章中介绍的机制基本相同,不同之处在于它应用于图像像素而非文本标记。
跳跃连接的使用以及模型的庞大规模导致去噪 U-Net 中存在冗余的特征提取,确保在去噪过程中不会丢失任何重要特征。然而,模型庞大的规模也使得识别相关特征变得复杂,有如大海捞针。注意力机制赋予模型聚焦于重要特征、忽略无关信息的能力,从而提升学习过程的有效性。
15.1.3 去噪 U-Net 模型的训练蓝图
去噪 U-Net 的输出是注入到带噪图像中的噪声。该模型通过最小化输出(预测噪声)与真实噪声(真实值)之间的差异来进行训练。
去噪 U-Net 模型利用了 U-Net 架构在捕捉局部和全局上下文方面的能力,使其在去除噪声的同时能有效保留重要细节,如边缘和纹理。这类模型广泛应用于多种领域,包括医学图像去噪、摄影图像修复等。图 15.3 展示了我们去噪 U-Net 模型的训练流程示意图。
图 15.3 去噪 U-Net 模型的训练过程示意图。我们首先获取干净的花卉图像作为训练集。然后在干净的花卉图像上添加噪声,将带噪图像输入到 U-Net 模型中。模型预测带噪图像中的噪声。我们将模型预测的噪声与实际注入花卉图像的噪声进行比较,并调整模型权重以最小化平均绝对误差。
第一步是收集花卉图像数据集。我们将使用牛津 102 花卉数据集作为训练集。所有图像会被调整为固定分辨率 64 × 64 像素,并将像素值归一化到 [–1, 1] 范围内。为了去噪,我们需要干净图像和带噪图像的配对。我们会根据公式 15.3,人工在干净的花卉图像中添加噪声,生成对应的带噪图像(图 15.3 中的步骤 2)。
接下来,我们构建结构如图 15.2 所示的去噪 U-Net 模型。在训练的每个 epoch 中,我们按批次遍历数据集。我们在花卉图像上添加噪声,将带噪图像及其对应的时间步长 t 一起输入到 U-Net 模型(步骤 3)。基于当前模型参数,U-Net 模型预测带噪图像中的噪声(步骤 4)。
然后,我们将预测的噪声与真实噪声进行比较,并在像素级别计算 L1 损失(即平均绝对误差)(步骤 5)。L1 损失通常比 L2 损失(均方误差)对异常值的敏感度更低,因此在此类任务中更受青睐。随后,我们调整模型参数以最小化 L1 损失(步骤 6),以便模型在下一轮迭代中做出更准确的预测。我们不断重复此过程,直到模型参数收敛。
15.2 准备训练数据
我们将使用牛津102花卉数据集作为训练数据,该数据集可在Hugging Face免费获取。该数据集包含约8000张花卉图像,可以通过之前安装的datasets库直接下载。
为了节省篇幅,我们会将大部分辅助函数和类放入两个本地模块:ch15util.py和unet_util.py。请从本书的GitHub仓库(github.com/markhliu/DG…)下载这两个文件,并将它们放置在你电脑的/utils/文件夹下。本章中的Python程序改编自Hugging Face的GitHub仓库(github.com/huggingface…)和Filip Basara的GitHub仓库(github.com/filipbasara…)。
你将使用Python将数据集下载到你的电脑。随后,我们会演示前向扩散过程,即逐步向训练数据中的干净图像添加噪声,直到它们变成随机噪声。最后,将训练数据整理成批次,方便后续训练去噪U-Net模型。
本章会用到以下Python库:datasets、einops、diffusers和openai。可在Jupyter Notebook中新建代码单元,执行以下命令安装:
!pip install datasets einops diffusers openai
然后按照屏幕提示完成安装。
15.2.1 以花卉图像作为训练数据
之前安装的datasets库中的load_dataset()方法可直接从Hugging Face下载牛津102花卉数据集。我们将使用matplotlib库展示部分花卉图像,以便直观了解训练集中的图像内容。
请在Jupyter Notebook中新建代码单元,运行如下代码:
from datasets import load_dataset
from utils.ch15util import transforms
dataset = load_dataset("huggan/flowers-102-categories", split="train") # ①
dataset.set_transform(transforms)
import matplotlib.pyplot as plt
from torchvision.utils import make_grid
# 将第一批16张图像绘制为网格
grid = make_grid(dataset[:16]["input"], 8, 2) # ②
plt.figure(figsize=(8,2), dpi=300)
plt.imshow(grid.numpy().transpose((1,2,0)))
plt.axis("off")
plt.show()
① 从Hugging Face下载图像
② 绘制前16张图像
运行上述代码后,你将看到数据集中前16张花卉图像,如图15.4所示。这些都是高分辨率的彩色花卉图像,我们已将每张图像大小统一调整为(3, 64, 64)。
我们将数据集按批次大小为4进行分组,以便后续用于训练去噪U-Net模型。选择批次大小为4是为了在训练时保持显存足够小,适合GPU使用。如果你的GPU显存较小,可以将批次大小调整为2甚至1:
import torch
resolution = 64
batch_size = 4
train_dataloader = torch.utils.data.DataLoader(
dataset, batch_size=batch_size, shuffle=True)
接下来,我们将编写代码并可视化前向扩散过程。
15.2.2 可视化前向扩散过程
我们在刚下载的本地模块 ch15util.py 中定义了 DDIMScheduler() 类。你可以打开该文件查看该类的定义,我们会用它来向图像中添加噪声。同时,我们还会结合训练好的去噪U-Net模型来生成干净图像。DDIMScheduler() 类负责管理每一步的时间步长和去噪序列,使模型能够进行确定性推理,从而通过去噪过程生成高质量样本。
首先,我们从训练集中选取四张干净图像,并生成与这些图像形状相同的噪声张量:
clean_images = next(iter(train_dataloader))["input"] * 2 - 1 # ①
print(clean_images.shape)
nums = clean_images.shape[0]
noise = torch.randn(clean_images.shape) # ②
print(noise.shape)
① 获取四张干净图像
② 生成一个与干净图像形状相同的张量 noise,其中每个值独立服从标准正态分布
上面代码的输出结果为:
torch.Size([4, 3, 64, 64])
torch.Size([4, 3, 64, 64])
图像和噪声张量的形状均为 (4, 3, 64, 64),表示批次中有4张图像,每张图像有3个颜色通道,高度和宽度均为64像素。
在前向扩散过程中,干净图像(x0,详见第一节)和随机噪声(xT)之间存在999个过渡的噪声图像。这些过渡图像是干净图像和噪声的加权和。随着时间步 t 从0增加到1000,干净图像的权重逐渐减小,而噪声的权重逐渐增加,具体见公式15.3。
接下来,我们生成并可视化一些过渡噪声图像。
from utils.ch15util import DDIMScheduler
noise_scheduler = DDIMScheduler(num_train_timesteps=1000) # ①
allimgs = clean_images
for step in range(200, 1001, 200): # ②
timesteps = torch.tensor([step - 1] * 4).long()
noisy_images = noise_scheduler.add_noise(clean_images,
noise, timesteps) # ③
allimgs = torch.cat((allimgs, noisy_images)) # ④
import torchvision
imgs = torchvision.utils.make_grid(allimgs, 4, 6)
fig = plt.figure(dpi=300)
plt.imshow((imgs.permute(2, 1, 0) + 1) / 2) # ⑤
plt.axis("off")
plt.show()
① 使用1000个时间步实例化 DDIMScheduler() 类
② 选择时间步为200、400、600、800和1000
③ 在这些时间步生成噪声图像
④ 将噪声图像与干净图像拼接
⑤ 显示所有图像
DDIMScheduler() 类中的 add_noise() 方法接收三个参数:干净图像、噪声张量和时间步,输出干净图像与噪声的加权和即噪声图像。权重随着时间步t的增加,干净图像的权重降低,噪声权重增加。如果你运行上述代码,会得到类似图15.5的图像效果。
图15.5 前向扩散过程示意图。第一列的四张图像是训练数据集中的干净图像。然后,我们从第1个时间步开始逐渐向这些图像中添加噪声,直到第1000个时间步。随着时间步的增加,图像中的噪声越来越多。第二列的四张图像是经过200个时间步后的噪声图像。第三列是400个时间步后的图像,噪声比第二列更多。最后一列是1000个时间步后的图像,完全变成随机噪声。
第一列显示的是四张没有噪声的干净图像。向右移动时,我们逐渐向图像添加越来越多的噪声。最右侧的列显示的则是纯随机噪声。
15.3 构建去噪U-Net模型
本章前面讨论了去噪U-Net模型的架构,这一节将指导你用Python和PyTorch实现它。
我们将构建的U-Net模型规模较大,包含超过1.33亿个参数,反映了其任务的复杂性。该模型通过对输入图像进行下采样和上采样,捕获图像的局部和全局特征。它使用多个卷积层,并通过跳跃连接将网络各层的特征融合,有助于保留空间信息,提升学习效果。
由于去噪U-Net模型规模庞大且具有冗余的特征提取,模型引入了缩放点积注意力机制(SDPA),使模型能够关注输入中最相关的部分。为了计算SDPA注意力,我们将图像展开,将像素视为序列,然后利用SDPA学习图像中不同像素之间的依赖关系,这与我们在第9章中学习文本中不同token之间的依赖类似。
15.3.1 去噪U-Net模型中的注意力机制
为了实现注意力机制,我们在本地模块 ch15util.py 中定义了一个 Attention() 类,代码如下:
import torch
from torch import nn, einsum
from einops import rearrange
class Attention(nn.Module):
def __init__(self, dim, heads=4, dim_head=32):
super().__init__()
self.scale = dim_head ** -0.5
self.heads = heads
hidden_dim = dim_head * heads
self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)
self.to_out = nn.Conv2d(hidden_dim, dim, 1)
def forward(self, x):
b, c, h, w = x.shape
qkv = self.to_qkv(x).chunk(3, dim=1) # ①
q, k, v = map(
lambda t: rearrange(t, 'b (h c) x y -> b h c (x y)', h=self.heads),
qkv) # ②
q = q * self.scale
sim = einsum('b h d i, b h d j -> b h i j', q, k)
attn = sim.softmax(dim=-1) # ③
out = einsum('b h i j, b h d j -> b h i d', attn, v) # ④
out = rearrange(out, 'b h (x y) d -> b (h d) x y', x=h, y=w)
return self.to_out(out) # ⑤
attn = Attention(128)
x = torch.rand(1, 128, 64, 64)
out = attn(x)
print(out.shape)
注释说明:
① 输入经过三个卷积层以生成查询(Q)、键(K)和值(V)
② 将Q、K、V拆分为多个并行头
③ 计算注意力权重
④ 计算每个头的注意力向量
⑤ 将多个头的向量合并为一个输出
运行结果:
torch.Size([1, 128, 64, 64])
这里使用的缩放点积注意力(SDPA)与第9章中对文本序列中token应用的相同,不同的是这里应用于图像像素。我们将图像像素展平为序列,利用SDPA提取图像中不同区域间的依赖,提升去噪效率。
代码示例中,创建了一个形状为 (1, 128, 64, 64) 的假设图像张量,表示批次大小1,128个特征通道,每个通道64×64像素。输入通过注意力层处理时,每个特征通道被展平为长度为4096(64×64)的序列。然后通过三个独立网络层生成查询Q、键K、值V,接着拆分为4个头。接下来计算每个头的注意力向量,步骤如下:
表示键向量 K 的维度。四个注意力头产生的注意力向量会被拼接回一个整体的注意力向量。
15.3.2 去噪 U-Net 模型
在你刚下载的本地模块 unet_util.py 中,定义了一个 UNet() 类,用来表示去噪 U-Net 模型。你可以查看文件中的定义,下面我会简要说明其工作原理。以下代码展示了 UNet() 类的一部分:
class UNet(nn.Module):
...
def forward(self, sample, timesteps): ①
if not torch.is_tensor(timesteps):
timesteps = torch.tensor([timesteps],
dtype=torch.long,
device=sample.device)
timesteps = torch.flatten(timesteps)
timesteps = timesteps.broadcast_to(sample.shape[0])
t_emb = sinusoidal_embedding(timesteps, self.hidden_dims[0])
t_emb = self.time_embedding(t_emb) ②
x = self.init_conv(sample)
r = x.clone()
skips = []
for block1, block2, attn, downsample in self.down_blocks: ③
x = block1(x, t_emb)
skips.append(x)
x = block2(x, t_emb)
x = attn(x)
skips.append(x)
x = downsample(x)
x = self.mid_block1(x, t_emb)
x = self.mid_attn(x)
x = self.mid_block2(x, t_emb) ④
for block1, block2, attn, upsample in self.up_blocks:
x = torch.cat((x, skips.pop()), dim=1) ⑤
x = block1(x, t_emb)
x = torch.cat((x, skips.pop()), dim=1)
x = block2(x, t_emb)
x = attn(x)
x = upsample(x)
x = self.out_block(torch.cat((x, r), dim=1), t_emb)
out = self.conv_out(x)
return {"sample": out} ⑥
注释说明:
① 模型接受一批带噪声的图像和对应的时间步作为输入。
② 嵌入后的时间步在模型的不同阶段被加入到图像特征中。
③ 输入经过收缩路径(编码器)。
④ 输入经过瓶颈层。
⑤ 输入经过扩展路径(解码器),并结合跳跃连接的特征。
⑥ 输出为输入图像中预测的噪声。
去噪 U-Net 的任务是基于时间步信息,预测输入带噪图像中的噪声。正如公式15.3所述,任意时间步 t 的带噪图像 可以表示为干净图像 与标准正态分布噪声 的加权和。随着时间步 t 从0到 T 递增,干净图像的权重逐渐减少,噪声的权重逐渐增加。因此,为了推断带噪图像中的噪声,去噪 U-Net 需要知道该图像处于哪个时间步。
时间步的嵌入使用类似 Transformer 中位置编码的正余弦函数(第9、10章讨论),生成一个128维的向量。随后,这些时间步嵌入会被扩展(broadcast)成与模型中不同层次图像特征相匹配的维度。例如,在第一个下采样块中,时间步嵌入被广播到形状 (128, 64, 64),并加到同尺寸的图像特征上。
接下来,我们通过实例化本地模块中的 UNet() 类创建去噪 U-Net 模型:
from utils.unet_util import UNet
device = "cuda" if torch.cuda.is_available() else "cpu"
resolution = 64
model = UNet(3, hidden_dims=[128, 256, 512, 1024],
image_size=resolution).to(device)
num = sum(p.numel() for p in model.parameters())
print("number of parameters: %.2fM" % (num / 1e6,))
print(model)
输出为:
number of parameters: 133.42M
如你所见,该模型拥有超过1.33亿个参数。由于参数众多,本章的训练过程耗时较长,大约需要3至4小时的GPU训练时间。但如果你没有GPU训练资源,训练好的权重文件也会在后续章节中提供下载链接。
15.4 训练与使用去噪 U-Net 模型
现在我们已经准备好了训练数据和去噪 U-Net 模型,接下来就可以使用训练数据对模型进行训练了。
在每个训练周期(epoch)中,我们会遍历训练数据中的所有批次。对于每张图像,我们会随机选取一个时间步(time step),并根据该时间步的值向干净的训练图像中添加噪声,得到一张带噪声的图像。然后将这些带噪图像及对应的时间步值输入到去噪 U-Net 模型中,模型的任务是预测图像中的噪声。我们将模型预测的噪声与真实噪声(即添加到图像中的实际噪声)进行比较,通过调整模型参数,最小化预测噪声与真实噪声之间的平均绝对误差(MAE)。
训练完成后,我们将使用训练好的模型生成花朵图像。生成过程共进行 50 步推理(即时间步值依次为 980、960、…、20 和 0)。起始时,从随机噪声开始,输入到训练好的模型中得到一张带噪声的图像,然后将这张图像再次输入模型进行去噪。重复此过程共 50 步,最终生成的图像将与训练集中的花朵图像无法区分。
15.4.1 训练去噪 U-Net 模型
接下来,我们先定义训练过程中的优化器(optimizer)和学习率调度器(learning rate scheduler)。
我们将使用 AdamW 优化器,这是 Adam 优化器的一个变体,整本书中都在使用。AdamW 优化器由 Ilya Loshchilov 和 Frank Hutter 提出,它将权重衰减(weight decay,一种正则化方式)从梯度更新步骤中解耦出来。具体来说,AdamW 不是直接对梯度施加权重衰减,而是在每次优化步骤后,直接对模型参数(权重)施加权重衰减。这样的修改有助于提高模型的泛化能力,避免权重衰减率随学习率同步调整。感兴趣的读者可以查阅 Loshchilov 和 Hutter 的原始论文深入了解 AdamW 优化器。
此外,我们还将使用 diffusers 库中的学习率调度器来动态调整训练过程中的学习率。训练初期使用较大学习率,有助于模型跳出局部最优;训练后期逐渐降低学习率,可以帮助模型更稳健准确地收敛到全局最优。学习率调度器定义如下所示。
from diffusers.optimization import get_scheduler
num_epochs = 100 # ①
optimizer = torch.optim.AdamW(
model.parameters(),
lr=0.0001,
betas=(0.95, 0.999),
weight_decay=0.00001,
eps=1e-8
) # ②
lr_scheduler = get_scheduler( # ③
"cosine",
optimizer=optimizer,
num_warmup_steps=300,
num_training_steps=(len(train_dataloader) * num_epochs)
)
① 模型训练总共进行 100 个周期
② 使用 AdamW 优化器
③ 使用 diffusers 库中的学习率调度器控制学习率
get_scheduler() 函数的具体实现可见 Hugging Face 的 GitHub: mng.bz/ZVo5 。
在前 300 个训练步骤(warmup steps)内,学习率线性从 0 增加到 0.0001(即 AdamW 优化器中设置的学习率);之后学习率按余弦函数从 0.0001 逐步减小到 0。
下面是训练去噪 U-Net 模型的代码示例:
for epoch in range(num_epochs):
model.train()
tloss = 0
print(f"start epoch {epoch}")
for step, batch in enumerate(train_dataloader):
clean_images = batch["input"].to(device) * 2 - 1
nums = clean_images.shape[0]
noise = torch.randn(clean_images.shape).to(device)
timesteps = torch.randint(
0,
noise_scheduler.num_train_timesteps,
(nums, ),
device=device
).long()
noisy_images = noise_scheduler.add_noise(clean_images, noise, timesteps) # ①
noise_pred = model(noisy_images, timesteps)["sample"] # ②
loss = torch.nn.functional.l1_loss(noise_pred, noise) # ③
loss.backward()
optimizer.step() # ④
lr_scheduler.step()
optimizer.zero_grad()
tloss += loss.detach().item()
if step % 100 == 0:
print(f"step {step}, average loss {tloss / (step + 1)}")
torch.save(model.state_dict(), 'files/diffusion.pth')
① 向训练集中的干净图像添加噪声
② 使用去噪 U-Net 预测噪声
③ 计算预测噪声与实际噪声的平均绝对误差(L1 损失)
④ 更新模型参数以最小化误差
在每个训练周期中,我们遍历训练集中的所有干净花朵图像批次,向它们添加噪声后输入去噪 U-Net,模型预测噪声,再根据预测噪声与真实噪声之间的差异调整模型参数,使平均绝对误差最小化。
本训练过程在 GPU 环境下运行,通常需要数小时完成。训练完成后,模型权重会保存在本地电脑。你也可以从我的网站 mng.bz/RNlD 下载预训练权重,下载后解压即可使用。
15.4.2 使用训练好的模型生成花朵图像
为了生成花朵图像,我们将使用 50 步推理。这意味着我们会在 t = 0 到 t = T 之间均匀选取 50 个时间步,且这里的 T = 1000。因此,50 个推理时间步为 t = 980、960、940、…、20 和 0。生成过程从纯随机噪声开始,对应于 t = 1000 时的图像。我们使用训练好的去噪 U-Net 模型对该噪声图像去噪,得到 t = 980 时的带噪图像。然后将 t = 980 的带噪图像输入模型继续去噪,得到 t = 960 的带噪图像。如此反复迭代,直到得到 t = 0 时的图像,也就是干净的最终图像。该过程通过本地模块 ch15util.py 中 DDIMScheduler() 类的 generate() 方法实现。
@torch.no_grad()
def generate(self, model, device, batch_size=1, generator=None,
eta=1.0, use_clipped_model_output=True, num_inference_steps=50):
imgs = []
image = torch.randn(
(batch_size, model.in_channels, model.sample_size, model.sample_size),
generator=generator
).to(device) # ①
self.set_timesteps(num_inference_steps) # ②
for t in tqdm(self.timesteps):
model_output = model(image, t)["sample"] # ③
image = self.step(model_output, t, image, eta,
use_clipped_model_output=use_clipped_model_output) # ④
img = unnormalize_to_zero_to_one(image)
img = img.cpu().permute(0, 2, 3, 1).numpy()
imgs.append(img) # ⑤
image = unnormalize_to_zero_to_one(image)
image = image.cpu().permute(0, 2, 3, 1).numpy()
return {"sample": image}, imgs
① 以随机噪声作为起始图像(即 t = 1000 时的图像)
② 使用 50 个推理时间步(t = 980、960、940、…、20、0)
③ 使用训练好的去噪 U-Net 模型预测噪声
④ 根据预测噪声生成图像
⑤ 将每个时间步的中间图像保存到列表 imgs 中
在该 generate() 方法中,我们创建了一个 imgs 列表,用于存储 t = 980、960、…、20、0 等时间步对应的所有中间图像,便于后续可视化去噪过程。generate() 方法返回一个包含最终生成图像和中间图像列表的字典。
接下来,我们使用前述 generate() 方法生成 10 张干净图像。
sd = torch.load('files/diffusion.pth', map_location=device)
model.load_state_dict(sd)
with torch.no_grad():
generator = torch.manual_seed(1) # ①
generated_images, imgs = noise_scheduler.generate(
model, device,
num_inference_steps=50,
generator=generator,
eta=1.0,
use_clipped_model_output=True,
batch_size=10) # ②
imgnp = generated_images["sample"]
import matplotlib.pyplot as plt
plt.figure(figsize=(10, 4), dpi=300)
for i in range(10): # ③
ax = plt.subplot(2, 5, i + 1)
plt.imshow(imgnp[i])
plt.xticks([])
plt.yticks([])
plt.tight_layout()
plt.show()
① 设置随机种子为 1,保证结果可复现
② 使用 generate() 方法生成 10 张干净图像
③ 将生成的图像绘制为 2×5 的网格图
我们将随机种子设置为 1,因此如果你使用我网站上的训练模型,能获得与图 15.6 中完全相同的生成效果。我们用 50 步推理调用前面定义的 generate() 方法生成了 10 张干净图像,最后以 2 行 5 列的形式展示这 10 张图像,效果如图 15.6 所示。
正如图15.6所示,生成的花朵图像看起来真实,且与训练数据集中的花朵十分相似。
练习15.1
修改代码清单15.8,将随机种子改为2,其余代码保持不变。重新运行代码,观察生成的图像效果。
generate() 方法还会返回一个列表 imgs,包含了50个中间推理步骤的所有图像。我们将利用这些图像来可视化去噪过程。
steps = imgs[9::10] # ①
imgs20 = []
for j in [1, 3, 6, 9]:
for i in range(5):
imgs20.append(steps[i][j]) # ②
plt.figure(figsize=(10, 8), dpi=300)
for i in range(20): # ③
k = i % 5
ax = plt.subplot(4, 5, i + 1)
plt.imshow(imgs20[i])
plt.xticks([])
plt.yticks([])
plt.tight_layout()
plt.title(f't={800-200*k}', fontsize=15, c="r")
plt.show()
① 保留时间步 t = 800、600、400、200 和 0 对应的图像
② 从10张花中选择4组(图15.6中的第2、第4、第7、第10张)
③ 将20张图像绘制成4×5的网格
列表 imgs 包含了10组花朵图像在全部50个推理时间步(t = 980、960、…、20、0)对应的图像,因此共有500张图像。我们从中选取5个时间步(t = 800、600、400、200和0)以及4朵不同的花(图15.6中的第2、第4、第7和第10张),并将这20张图像以4行5列的形式绘制出来,效果如图15.7所示。
图15.7 展示了训练好的去噪 U-Net 模型如何逐步将随机噪声转化为干净的花朵图像。我们先将随机噪声输入训练模型,得到时间步 t = 980 的图像;再将 t = 980 的带噪图像输入模型,得到 t = 960 的图像。我们重复该过程共 50 步推理,直到获得 t = 0 的图像。图中第一列显示了 t = 800 时的四朵花;第二列显示同样四朵花在 t = 600 时的状态……最后一列展示了 t = 0 时的四朵花(即干净的花朵图像)。
图15.7 第一列中,四朵花在 t = 800 时的图像接近随机噪声;第二列显示 t = 600 时的花朵,开始呈现花的形态。向右移动,图像逐渐变得更加清晰,最右一列为 t = 0 时的四张清晰花朵图像。
了解了扩散模型的工作原理后,我们接下来讨论文本到图像的生成。文本到图像的 Transformer 模型(如 DALL·E 2、Imagen 和 Stable Diffusion)生成图像的过程,与前面章节中讨论的逆扩散过程非常相似,只不过这些模型在生成图像时,会将文本的嵌入向量作为条件信号输入。
15.5 文本到图像的 Transformer
OpenAI 的 DALL·E 2、Google 的 Imagen 和 Stability AI 的 Stable Diffusion 等文本到图像 Transformer,均使用扩散模型从文本描述生成图像。这类模型的核心组件是扩散模型。文本生成图像的过程包括将文本输入编码为潜在表示,并将其作为扩散模型的条件信号。这些 Transformer 通过引导编码后的文本,迭代地对随机噪声向量去噪,从而生成与文本描述相符的逼真图像。
所有文本到图像 Transformer 的关键,在于模型能理解不同模态的内容——即理解文本描述并将其关联到图像,反之亦然。
本节以 OpenAI 的 CLIP 模型为例。CLIP 是 DALL·E 2 的关键组件。我们将介绍 CLIP 如何训练以理解文本描述与图像之间的关联,随后用一个简短的 Python 程序展示如何使用 DALL·E 2 根据文本提示生成图像。
15.5.1 CLIP:一种多模态 Transformer
近年来,计算机视觉与自然语言处理(NLP)的交叉领域取得了显著进展,其中之一便是 OpenAI 提出的 CLIP 模型。该创新模型旨在基于自然语言理解和解释图像,这种能力在图像生成和图像分类等多种应用中具有巨大潜力。
CLIP 是一种多模态 Transformer,连接视觉数据和文本数据之间的鸿沟。它通过关联图像与相应的文本描述来训练模型理解图像。与传统需要对图像进行明确标注的模型不同,CLIP 利用庞大的图像与自然语言描述数据集,学习对视觉概念更具通用性的表示。
图15.8 展示了 OpenAI 的 CLIP 模型是如何训练的。首先收集了一个大规模的文本-图像对训练数据集。模型中的文本编码器将文本描述压缩为一个维度为 D 的文本嵌入向量,图像编码器则将对应的图像转换为同样维度为 D 的图像嵌入向量。在训练过程中,一批 N 对文本-图像对被转换为 N 个文本嵌入和 N 个图像嵌入。CLIP 采用对比学习方法,最大化匹配文本-图像对嵌入之间的相似度(图中对角线元素之和),同时最小化非匹配文本-图像对嵌入之间的相似度(图中非对角线元素之和)。
CLIP 模型的训练过程如图15.8所示,首先收集一个包含大量图像及其对应文本描述的大规模数据集。OpenAI 利用多个来源的数据,包括公开数据集和网络爬取数据,确保视觉和文本内容的多样性。数据集随后经过预处理,标准化图像尺寸使其统一,并对文本进行分词处理,为模型输入做好准备。
CLIP 采用双编码器架构,包括图像编码器和文本编码器。图像编码器负责处理输入图像,文本编码器处理对应的文本描述。这两个编码器将图像和文本映射到一个共享的嵌入空间,使其能够进行比较和对齐。
CLIP 训练的核心是对比学习方法。对于数据集中每一批 N 对图文,模型旨在最大化匹配文本-图像对嵌入之间的相似度(如图15.8中对角线元素之和所示),同时最小化非匹配文本-图像对嵌入之间的相似度(非对角线元素之和)。图15.9展示了文本到图像 Transformer(如 DALL·E 2)如何基于文本提示生成逼真图像的示意图。
图15.9 展示了文本到图像 Transformer(如 DALL·E 2)如何根据文本提示生成图像。训练好的文本到图像 Transformer 中的文本编码器首先将提示中的文本描述转换为文本嵌入向量。该文本嵌入向量输入 CLIP 模型,获得表示图像潜在空间的先验向量。文本嵌入与先验向量拼接为一个条件向量。为了生成图像,U-Net 去噪器首先以随机噪声向量为输入,结合条件向量生成一张带噪图像。随后,U-Net 以该带噪图像和条件向量为输入,生成噪声更少的图像。此过程重复多次,直到最终生成一张干净的图像。
文本到图像 Transformer 的图像生成过程类似于本章前面讨论的逆扩散过程。以 OpenAI 研究人员于 2022 年提出的 DALL·E 2 为例:模型中的文本编码器先将提示文本转换为文本嵌入,随后将其输入 CLIP 模型得到潜在空间中的先验向量,二者拼接形成条件向量。第一步,向 U-Net 去噪器输入随机噪声向量,结合条件向量生成带噪图像;第二步,将上一步生成的带噪图像再次输入 U-Net 去噪器,结合条件向量生成新的带噪图像。该过程重复多次,最终输出为一张干净的图像。
15.5.2 使用 DALL·E 2 进行文本到图像生成
了解了文本到图像 Transformer 的工作原理后,我们编写一个 Python 程序,使用 DALL·E 2 根据文本提示生成图像。
首先,你需要申请一个 OpenAI API 密钥。OpenAI 提供不同的定价方案,费用取决于处理的 token 数量和使用的模型类型。访问 chat.openai.com/auth/login 并点击“注册”按钮创建账号。登录后,进入 platform.openai.com/api-keys 查看你的 API 密钥,并妥善保存以备后续使用。我们可以用 OpenAI 的 DALL·E 2 生成图像。
from openai import OpenAI
openai_api_key = "你的实际 OpenAI API 密钥" # ①
client = OpenAI(api_key=openai_api_key) # ②
response = client.images.generate(
model="dall-e-2",
prompt="an astronaut in a space suit riding a unicorn",
size="512x512",
quality="standard",
n=1,
) # ③
image_url = response.data[0].url
print(image_url) # ④
① 确保在此处以字符串形式填写你真实的 OpenAI API 密钥
② 实例化 OpenAI() 类,创建客户端代理
③ 使用 images.generate() 方法,根据文本提示生成图像
④ 输出生成图像的 URL
你需要将之前获得的 OpenAI API 密钥填入代码清单15.10中。我们通过实例化 OpenAI() 类创建代理。生成图像时,需要指定模型、文本提示和图像大小。此处示例使用了提示“an astronaut in a space suit riding a unicorn”。代码返回一个 URL,供我们查看和下载生成的图像。该 URL 有效期为一小时。生成的图像示例见图15.10。
图15.10 展示了使用文本提示“an astronaut in a space suit riding a unicorn”(穿宇航服的宇航员骑独角兽)由 DALL·E 2 生成的一张图像。
请你亲自运行代码清单15.10,看看 DALL·E 2 为你生成了怎样的图像。需要注意的是,由于 DALL·E 2(以及所有大型语言模型)的输出是随机的,而非确定性的,因此你的结果可能会有所不同。
练习15.2
申请一个 OpenAI API 密钥,然后修改代码清单15.10,将文本提示改为“a cat in a suit working on a computer”(穿西装的猫在电脑前工作),生成相应图像。
本章你学习了基于扩散模型的工作原理及其在文本到图像 Transformer(如 OpenAI 的 CLIP 模型)中的重要作用。你还了解了如何获取 OpenAI API 密钥,并通过一个简短的 Python 脚本,利用集成了 CLIP 的 DALL·E 2,根据文本描述生成图像。
下一章,你将继续使用之前获得的 OpenAI API 密钥,调用预训练大型语言模型(LLM)生成丰富多样的内容,包括文本、音频和图像。同时,你将学习如何将 LangChain Python 库与其他 API 集成,实现一个无所不知的个人智能助理。
总结
在正向扩散过程中,我们逐步向干净图像中添加少量随机噪声,直到图像变成纯噪声。相反,在逆向扩散过程中,我们从随机噪声开始,利用去噪模型逐步去除图像中的噪声,将噪声恢复为干净的图像。
U-Net 架构最初设计用于生物医学图像分割,具有对称结构,包括一个收缩的编码器路径和一个扩展的解码器路径,中间通过瓶颈层相连。用于去噪时,U-Net 经过改造以在去除噪声的同时保留细节。跳跃连接(skip connection)连接编码器和解码器中相同空间尺寸的特征图,有助于保留编码过程中的下采样可能丢失的空间信息,如边缘和纹理。
在去噪 U-Net 中引入注意力机制,使模型能聚焦于重要特征,忽略无关信息。通过将图像像素视为序列,注意力机制学习像素之间的依赖关系,类似于自然语言处理中学习词元依赖的方式,从而提升模型识别相关特征的能力。
OpenAI 的 DALL·E 2、Google 的 Imagen 以及 Stability AI 的 Stable Diffusion 等文本到图像 Transformer,利用扩散模型根据文本描述生成图像。它们先将文本编码为潜在表示,作为扩散模型的条件输入,然后引导扩散模型迭代去噪随机噪声向量,生成符合文本描述的逼真图像。