🔎大家好,我是Sonhhxg_柒,希望你看完之后,能对你有所帮助,不足请指正!共同学习交流🔎
📝个人主页-Sonhhxg_柒的博客_CSDN博客 📃
🎁欢迎各位→点赞👍 + 收藏⭐️ + 留言📝
📣系列专栏 - 机器学习【ML】 自然语言处理【NLP】 深度学习【DL】
🖍foreword
✔说明⇢本人讲解主要包括Python、机器学习(ML)、深度学习(DL)、自然语言处理(NLP)等内容。
如果你对这个系列感兴趣的话,可以关注订阅哟👋
文章目录
2020 年底,一类鲜为人知的称为扩散模型的模型开始在机器学习领域引起轰动。研究人员弄清楚了如何使用这些模型生成质量高于以往技术的合成图像。随后出现了一系列论文,提出了改进和修改的建议,进一步提高了质量。到 2021 年底,像 GLIDE 这样的模型在文本到图像的任务上展示了令人难以置信的结果,几个月后,这些模型通过 DALL-E 2 和 Stable Diffusion 等工具进入了主流。这些模型使任何人都可以轻松地生成图像,只需输入他们想要看到的内容的文本描述即可。
在本文中,我们将深入探讨这些模型如何工作的细节。我们将概述使它们如此强大的关键见解,使用现有模型生成图像以感受它们的工作原理,然后训练我们自己的模型以进一步加深这种理解。该领域仍在快速发展,但此处涵盖的主题应该为您奠定坚实的基础。
关键见解:迭代优化
那么是什么让扩散模型如此强大呢?以前的技术,例如 VAE 或 GAN,通过模型的单个前向传递生成最终输出。这意味着模型必须在第一次尝试时就把所有事情都做对。如果它犯了一个错误,它就不能回头修复它。另一方面,扩散模型通过迭代许多步骤来生成它们的输出。这种“迭代优化”允许模型纠正之前步骤中的错误并逐渐改进输出。为了说明这一点,让我们看一个扩散模型的例子。
我们可以使用 Hugging Face diffusers 库加载预训练模型。管道可用于直接创建图像,但这并没有向我们展示幕后发生的事情:
-
# 加载pipeline
-
image_pipe
= DDPMPipeline.
from_pretrained(
"google/ddpm-celebahq-256")
-
image_pipe.
to(device);
-
-
# 一个图片实例
-
image_pipe().images[
0]
我们可以逐步重新创建采样过程,以更好地了解模型生成图像时发生的情况。我们用随机噪声初始化样本 x,然后在模型中运行它 30 个步骤。在右侧,您可以看到模型对最终图像在特定步骤的外观的预测 - 请注意,初始预测并不是特别好!我们没有直接跳到最终预测图像,而是仅在预测方向(如左侧所示)对 x 进行了少量修改。然后,我们在下一步中再次将这个新的、稍微好一点的 x 输入模型,希望得到稍微改进的预测,它可以用来更新 x 一点点,等等。通过足够的步骤,该模型可以生成一些令人印象深刻的逼真图像。
-
# 一批
4幅图像的随机起点
-
x
= torch.randn(
4,
3,
256,
256).
to(device)
-
-
# 将时间步长设置得更低
-
image_pipe.scheduler.
set_timesteps(num_inference_steps
=
30)
-
-
# 循环采样时间步长
-
for i, t
in enumerate(image_pipe.scheduler.timesteps):
-
-
# 在给定当前样本x和时间步长t的情况下获得预测
-
with torch.
no_grad():
-
noise_pred
= image_pipe.unet(x, t)[
"sample"]
-
-
# 使用调度程序计算更新后的样本应该是什么样子
-
scheduler_
output
= image_pipe.scheduler.step(noise_pred, t, x)
-
-
# 更新x
-
x
= scheduler_
output.prev_sample
-
-
# 偶尔显示x和预测的去噪图像
-
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_x
0
= scheduler_
output.pred_original_sample
-
grid
= torchvision.utils.make_grid(pred_x
0, 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()
如果该代码看起来有点吓人,请不要担心——我们将在本文中解释这一切是如何工作的。现在,只关注结果。
这种学习如何逐渐改进“损坏的”输入的核心思想可以应用于广泛的任务。在本章中,我们将重点关注无条件图像生成——即生成与训练数据相似的图像,而不对这些生成的样本的外观进行额外控制。扩散模型也被应用于音频、视频、文本等。虽然大多数实现都使用我们将在此处介绍的“去噪”方法的某些变体,但正在出现使用不同类型的“损坏”以及迭代优化的新方法,这可能会使该领域超越当前对去噪扩散的具体关注。激动人心的时刻!
训练扩散模型
在本节中,我们将从头开始训练扩散模型,以更好地了解它们的工作原理。我们将从使用 Hugging Face 扩散器库中的组件开始。随着本章的推进,我们将逐渐揭开每个组件如何工作的神秘面纱。与其他类型的生成模型相比,训练扩散模型相对简单。我们反复:
-
从训练数据中加载一些图像。
-
添加不同数量的噪声。请记住,我们希望模型能够很好地估计如何“修复”(去噪)极其嘈杂的图像和接近完美的图像。
-
将输入的噪声版本馈送到模型中。
-
评估模型在对这些输入进行去噪方面的表现。
-
使用此信息更新模型权重。
为了使用经过训练的模型生成新图像,我们从完全随机的输入开始,并重复将其输入模型,根据模型预测对每次迭代的输入进行少量更新。正如我们将看到的,有许多采样方法试图简化这个过程,以便我们可以用尽可能少的步骤生成好的图像。
数据
对于这个例子,我们将使用来自 Hugging Face Hub 的图像数据集——具体来说, 这个包含 1000 张蝴蝶图片的集合。稍后,在项目部分,您将看到如何使用您自己的数据。
dataset = load_dataset("huggan/smithsonian_butterflies_subset", split="train")
在使用这些数据训练模型之前,我们需要做一些准备工作。图像通常表示为“像素”网格,三个颜色通道(红色、绿色和蓝色)中的每一个的颜色值都在 0 到 255 之间。为了处理这些并使它们为训练做好准备,我们: - 将它们调整为固定大小 - (可选)通过随机水平翻转它们来添加一些增强,有效地将我们的数据集的大小加倍 - 将它们转换为 PyTorch 张量(代表颜色值在 0 和 1 之间浮动)- 将它们归一化为平均值为 0,值在 -1 和 1 之间
我们可以通过以下方式完成所有这些torchvision.transforms
:
-
image_
size
=
64
-
-
# 定义数据扩充
-
preprocess
= transforms.Compose(
-
[
-
transforms.Resize((image_
size, image_
size)), # Resize
-
transforms.RandomHorizontalFlip(), # 随机翻转(数据扩充)
-
transforms.ToTensor(), # Convert
to tensor (
0,
1)
-
transforms.Normalize([
0.5], [
0.5]), # Map
to (-
1,
1)
-
]
-
)
接下来,我们需要创建一个数据加载器来批量加载应用这些转换的数据:
-
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
-
)
我们可以通过加载单个批次并检查图像来检查这是否有效。
-
batch
=
next(iter(train_dataloader))
-
print(
'Shape:', batch[
'images'].shape,
-
'\nBounds:', batch[
'images'].min().item(),
'to', batch[
'images'].m
ax().item())
-
show_images(batch[
'images'][:
8]
*
0.5
+
0.5) # 注意:我们映射回(
0,
1)进行显示
-
Shape: torch.
Size([
32,
3,
64,
64])
-
Bounds: -
0.9921568632125854
to
1.0
添加噪音
我们如何逐渐破坏我们的数据?最常见的方法是在图像中添加噪声。我们添加的噪音量由噪音时间表控制。不同的论文和方法以不同的方式解决这个问题,我们将在本章后面探讨。现在,让我们看看基于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))
在训练期间,我们将随机选择时间步长。调度程序采用一些参数(beta_start 和 beta_end),用于确定给定时间步应存在多少噪声。我们将在第 X 节中更详细地介绍调度程序。
The UNet
UNet 是为图像分割等任务而发明的卷积神经网络,其中所需的输出与输入具有相同的空间范围。它由一系列减少输入空间大小的“下采样”层组成,然后是一系列再次增加输入空间范围的“上采样”层。下采样层通常还跟有一个“跳跃连接”,它将下采样层的输出连接到上采样层的输入。这允许上采样层“看到”网络早期的高分辨率表示,这对于具有类似图像输出的任务非常有用,在这些任务中,这种高分辨率信息特别有用。
扩散器库中使用的 UNet 架构比Ronneberger 等人在 2015 年提出的原始 UNet更先进,增加了注意力和残差块等功能。我们稍后会仔细研究,但这里的关键特征是它可以接受输入(噪声图像)并产生具有相同形状的预测(预测噪声)。对于扩散模型,UNet 通常还会将时间步作为附加条件,我们将在 UNet 深入探讨部分再次探讨这一点。
以下是我们如何创建一个 UNet 并通过它提供我们的一批噪声图像:
-
# Create a UNet
2DModel
-
model
= UNet
2DModel(
-
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
torch.Size([8, 3, 64, 64])
请注意,输出与输入的形状相同,这正是我们想要的。
训练
现在我们已经准备好模型和数据,可以训练它了。我们将使用学习率为 3e-4 的 AdamW 优化器。对于每个训练步骤,我们:
-
加载一批图像。
-
向图像添加噪声,选择随机时间步长来确定添加了多少噪声。
-
将噪声图像输入模型。
-
计算损失,这是模型预测和目标之间的均方误差——在本例中是 我们添加到图像中的噪声。这称为噪声或“epsilon”目标。您可以在第 X 节中找到有关不同培训目标的更多信息。
-
反向传播损失并使用优化器更新模型权重。
以下是代码中的所有内容:
-
num_epochs
=
50 # 我们应该在数据中运行多少次?
-
lr
=
1e-
4 # 我们应该使用什么学习率
-
model
= model.
to(device) # 我们正在训练的模型(在上一节中定义)
-
optimizer
= torch.optim.AdamW(model.parameters(), lr
=lr) # The optimizer
-
losses
= [] # 用于存储损失值以供以后绘图的位置
-
-
# 训练模型(这需要一段时间!)
-
for epoch
in range(num_epochs):
-
for step, batch
in enumerate(train_dataloader):
-
-
# 加载输入图像
-
clean_images
= batch[
"images"].
to(device)
-
-
# 采样噪声以添加到图像中
-
noise
= torch.randn(clean_images.shape).
to(clean_images.device)
-
-
# 为每个图像采样一个随机时间步长
-
timesteps
= torch.randint(
-
0,
-
scheduler.num_train_timesteps,
-
(clean_images.shape[
0],),
-
device
=clean_images.device,
-
).long()
-
-
# 根据时间步长向干净的图像添加噪波
-
noisy_images
= scheduler.
add_noise(clean_images, noise, timesteps)
-
-
# 获取噪声的模型预测
-
noise_pred
= model(noisy_images, timesteps,
return_dict
=
False)[
0]
-
-
# 将预测与实际噪声进行比较:
-
loss
= F.mse_loss(noise_pred, noise)
-
-
# 存储损失以备日后绘图
-
losses.append(loss.item())
-
-
# 基于此损失使用优化器更新模型参数
-
loss.backward(loss)
-
optimizer.step()
-
optimizer.
zero_grad()
在 GPU 上运行上述代码需要一个小时左右的时间,所以在等待或减少 epoch 数时喝点茶。这是训练后的损失曲线:
# Plot the loss curve:
plt.plot(losses);
随着模型学习对图像进行去噪,损失曲线呈下降趋势。由于根据每次迭代的时间步长随机采样向图像添加了不同数量的噪声,因此曲线相当嘈杂。仅通过查看噪声预测的均方误差很难判断该模型是否擅长生成样本,因此让我们继续下一节,看看它的表现如何。
采样
扩散器库使用“管道”的概念,它将使用扩散模型生成样本所需的所有组件捆绑在一起:
pipeline = DDPMPipeline(unet=model, scheduler=scheduler)
ims = pipeline(batch_size=4).images
show_images(ims, nrows=1)
当然,将创建样本的工作卸载到管道并不能真正向我们展示正在发生的事情。因此,这是一个简单的采样循环,显示了模型如何逐渐细化输入图像:
-
# 随机起点(
4张随机图像):
-
sample
= torch.randn(
4,
3,
64,
64).
to(device)
-
-
for i, t
in enumerate(scheduler.timesteps):
-
-
# 获得模型预测
-
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)
这与我们在本章开头用来说明迭代优化思想的代码相同,但希望现在您对这里发生的事情有了更好的理解。我们从一个完全随机的输入开始,然后由模型通过一系列步骤对其进行细化。每个步骤都是对输入的一个小更新,基于模型对该时间步的噪声的预测。我们仍在抽象出调用背后的一些复杂性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。冷扩散纸中使用的不同降解的图示
尽管如此,由于以下几个原因,添加噪音仍然是最受欢迎的方法:
-
我们可以轻松控制添加的噪声量,实现从“完美”到“完全损坏”的平滑过渡。对于降低图像分辨率之类的情况,情况并非如此,这可能会导致“离散”过渡。
-
我们可以有许多有效的随机起点用于推理,这与某些可能只有有限数量的可能初始(完全损坏)状态的方法不同,例如全黑图像或单像素图像。
因此,至少目前,我们将坚持添加噪音作为我们的腐败方法。接下来,让我们仔细看看如何为图像添加噪点。
从简单开始
我们有一些图像 (x),我们想以某种方式将它们与一些随机噪声结合起来。
-
x
=
next(iter(train_dataloader))[
'images'][:
8]
-
noise
= torch.rand_like(x)
我们可以做到这一点的一种方法是在它们之间线性插值 (lerp) 一定量。这给了我们一个函数,当“数量”从 0 到 1 变化时,它可以从原始图像 x 平滑过渡到纯噪声:
-
def corrupt(x, noise, amount):
-
amount
= amount.view(-
1,
1,
1,
1) # 确保它是可广播的
-
return x
*(
1-amount)
+ noise
*amount
让我们看看这对一批数据的影响,噪声量从 0 到 1 不等:
-
amount
= torch.linspace(
0,
1,
8)
-
noised_x
= corrupt(x, noise, amount)
-
show_images(noised_x
*
0.5
+
0.5)
这似乎正是我们想要的,从原始图像平滑过渡到纯噪声。现在,我们在这里创建了一个噪音时间表,它接受从 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)
现在我们有了一些可以直接与扩散器库中使用的调度程序进行比较的东西,例如我们在训练期间使用的 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))
数学
文献中有许多相互竞争的符号和方法。例如,一些论文在连续时间内参数化噪声时间表 ,其中 t 从 0(无噪声)到 1(完全损坏)——就像我们 corrupt
在上一节中的函数一样。其他人使用离散时间 方法,整数时间步长从 0 运行到某个大数 T,通常为 1000。可以像我们在课堂上所做的那样在这两种方法之间进行转换SimpleScheduler
——只需确保在比较不同模型时保持一致. 我们将在这里坚持使用离散时间方法。
深入研究数学的好起点是前面提到的 DDPM 论文。您可以在此处找到带注释的实现,这是理解此方法的重要附加资源。
该论文首先指定从时间步长 t-1 到时间步长 t 的单个噪声步长。他们是这样写的:
![](https://i-blog.csdnimg.cn/blog_migrate/39d932063198153edfaf237e673cc515.png)
这里为βt 所有时间步 t 定义,用于指定在每个步骤添加多少噪声。这个符号可能有点吓人,但这个等式告诉我们的是,噪声越大 𝐱t 是一个均值为 √¯1-βt · xt-1 和方差βt 换句话说,𝐱t 是混合的𝐱t-1(按比例缩放√1-βt ) 和一些随机噪声,我们可以将其视为按比例缩放的单位方差噪声t . 鉴于xt-1 和一些噪音∊ ,我们可以从这个分布中抽样得到xt 和:
为了在时间步长 t 获得噪声输入,我们可以从 t=0 开始并重复应用这个单一步骤,但这将是非常低效的。相反,我们可以找到一个公式来一次移动到任何时间步长 t。我们定义ɑ t =1 - βt 然后使用以下公式:
![](https://i-blog.csdnimg.cn/blog_migrate/2a047c186ccf08531b19a499168cad77.png)
在哪里 -∊ 是一些具有单位方差的高斯噪声 -ɑ¯('alpha_bar') 是所有的累积乘积ɑ值到时间t.
所以xt 是混合的 x0(按比例缩放 √¯ɑt) 和∊ (按比例缩放√1-¯ɑt). 在扩散器库中 ɑ¯值存储在 scheduler.alphas_cumprod
. 知道这一点,我们可以绘制原始图像的比例因子x0和噪音 ∊跨给定调度程序的不同时间步长:
plot_scheduler(DDPMScheduler()) # The default scheduler
我们上面的 SimpleScheduler 只是在原始图像和噪声之间线性混合,我们可以看到如果我们绘制比例因子(相当于 √¯ɑt和 (√1-¯ɑt)在 DDPM 案例中):
plot_scheduler(SimpleScheduler())
良好的噪音时间表将确保模型看到不同噪音水平的混合图像。最佳选择将根据训练数据而有所不同。可视化更多选项,请注意:
-
将 beta_end 设置得太低意味着我们永远不会完全擦除图像,因此模型永远不会看到像随机噪声这样的东西用作推理的起点。
-
将 beta_end 设置得非常高意味着大部分时间步都花在了几乎完全的噪声上,这将导致训练性能不佳。
-
不同的 beta 时间表给出不同的曲线。
“余弦”时间表是一种流行的选择,因为它提供了从原始图像到噪声的平滑过渡。
-
fig, (ax)
= plt.subplots(
1,
1, figsize
=(
8,
5))
-
plot_scheduler(DDPMScheduler(beta_schedule
=
"linear"),
-
label
=
'default schedule', ax
=ax, plot_both
=
False)
-
plot_scheduler(DDPMScheduler(beta_schedule
=
"squaredcos_cap_v2"),
-
label
=
'cosine schedule', ax
=ax, plot_both
=
False)
-
plot_scheduler(DDPMScheduler(beta_
end
=
0.003, beta_schedule
=
"linear"),
-
label
=
'Low beta_end', ax
=ax, plot_both
=
False)
-
plot_scheduler(DDPMScheduler(beta_
end
=
0.1, beta_schedule
=
"linear"),
-
label
=
'High beta_end', ax
=ax, plot_both
=
False)
此处显示的所有计划都称为“方差保留”(VP),这意味着模型输入的方差在整个计划中保持接近 1。您可能还会遇到“方差爆炸”(VE) 公式,其中只是将不同数量的噪声添加到原始图像(导致高方差输入)。我们将在抽样一章中对此进行更多讨论。我们的 SimpleScheduler 几乎是一个 VP 调度,但是由于线性插值,方差没有完全保留。
与许多与扩散相关的主题一样,不断有新论文探讨噪声时间表的主题,因此当您阅读本文时,可能会有大量可供尝试的选项!
输入分辨率和缩放比例的影响
直到最近才被忽视的噪声时间表的一个方面是输入大小和缩放的影响。许多论文在小规模数据集和低分辨率下测试潜在的调度器,然后使用性能最好的调度器在更大的图像上训练他们的最终模型。如果我们向两个不同大小的图像添加相同数量的噪声,就会发现这个问题。
图 1-2。比较不同尺寸图片加噪声的效果
高分辨率的图像往往包含大量冗余信息。这意味着即使单个像素被噪声遮挡,周围的像素也包含足够的信息来重建原始图像。对于低分辨率图像,情况并非如此,在低分辨率图像中,单个像素可以包含大量信息。这意味着与向高分辨率图像添加等量噪声相比,向低分辨率图像添加相同量的噪声会导致图像损坏得多。
两篇独立论文对这种效应进行了彻底研究,这两篇论文均于 2023 年 1 月发表。每篇论文都使用新的见解来训练能够生成高分辨率输出的模型,而无需任何以前必需的技巧。Hoogeboom 等人的简单扩散引入了一种根据输入大小调整噪声时间表的方法,允许针对新的目标分辨率适当修改针对低分辨率图像优化的时间表。一篇名为“关于扩散模型的噪声调度的重要性”的论文Ting Chen 进行了类似的实验,并注意到另一个关键变量:输入缩放。也就是说,我们如何表示我们的图像?如果图像表示为 0 和 1 之间的浮点数,那么它们的方差将低于噪声(通常是单位方差),因此对于给定的噪声水平,信噪比将低于图像表示为 -1 和 1 之间的浮点数(我们在上面的训练示例中使用过)或其他值。缩放输入图像会改变信噪比,因此修改此缩放比例是我们在对较大图像进行训练时可以调整的另一种方式。
深入:UNet 和替代方案
现在让我们讨论做出最重要预测的实际模型!回顾一下,该模型必须能够接收嘈杂的图像并估计如何对其进行降噪。这需要一个可以接收任意大小的图像并输出相同大小的图像的模型。此外,该模型应该能够在像素级别进行精确预测,同时还可以捕获有关整个图像的更高级别的信息。一种流行的方法是使用称为 UNet 的架构。UNets 于 2015 年发明用于医学图像分割,此后成为各种图像相关任务的流行选择。与我们在上一章中看到的自动编码器和 VAE 一样,UNet 由一系列“下采样”和“上采样”块组成。下采样块负责减小图像的大小,而上采样块负责增加图像的大小。下采样块通常由一系列卷积层组成,然后是池化或下采样层。上采样块通常由一系列卷积层组成,然后是上采样或“转置卷积”层。转置卷积层是一种特殊类型的卷积层,它增加图像的尺寸,而不是缩小图像。
常规自动编码器或 VAE 不是此任务的好选择的原因是它们在像素级别进行精确预测的能力较弱,因为输出必须完全从低维潜在空间重建。在 UNet 中,下采样和上采样块通过“跳过连接”连接,这允许信息直接从下采样块流向上采样块。这使得模型可以在像素级别进行精确预测,同时还可以捕获有关整个图像的更高级别的信息。
一个简单的 UNet
为了更好地理解 UNet 的结构,让我们从头开始构建一个简单的 UNet。
图 1-3。我们简单的 UNet 架构
这个 UNet 接受 32 像素分辨率的单通道输入,并以 32 像素分辨率输出单通道输出,我们可以用它来为 MNIST 数据集构建扩散模型。编码路径有三层,解码路径有三层。每层都包含一个卷积,然后是一个激活函数和一个上采样或下采样步骤(取决于我们是在编码还是解码路径中)。跳过连接允许信息直接从下采样块流向上采样块,并通过将下采样块的输出添加到相应上采样块的输入来实现。一些 UNet 将下采样块的输出连接到相应上采样块的输入,并且还可能在跳过连接中包含其他层。
-
from torch import nn
-
-
class BasicUNet(nn.Module):
-
""
"UNet的最低限度实施。"
""
-
def __init__(
self,
in_channels
=
1, out_channels
=
1):
-
super().__init__()
-
self.
down_layers
= torch.nn.ModuleList([
-
nn.Conv
2d(
in_channels,
32, kernel_
size
=
5, padding
=
2),
-
nn.Conv
2d(
32,
64, kernel_
size
=
5, padding
=
2),
-
nn.Conv
2d(
64,
64, kernel_
size
=
5, padding
=
2),
-
])
-
self.
up_layers
= torch.nn.ModuleList([
-
nn.Conv
2d(
64,
64, kernel_
size
=
5, padding
=
2),
-
nn.Conv
2d(
64,
32, kernel_
size
=
5, padding
=
2),
-
nn.Conv
2d(
32, out_channels, kernel_
size
=
5, padding
=
2),
-
])
-
self.act
= nn.SiLU() # 激活功能
-
self.downscale
= nn.MaxPool
2d(
2)
-
self.upscale
= nn.Upsample(scale_factor
=
2)
-
-
def forward(
self, x):
-
h
= []
-
for i, l
in enumerate(
self.
down_layers):
-
x
=
self.act(l(x)) # 通过层和激活功能
-
if i
<
2: # 对于除第三层(最后一层)外的所有下层:
-
h.append(x) # 存储跳过连接的输出
-
x
=
self.downscale(x) # 向下缩放为下一层做好准备
-
-
for i, l
in enumerate(
self.
up_layers):
-
if i
>
0: #
For
all except the
first
up layer
-
x
=
self.upscale(x) # Upscale
-
x
+
= h.pop() # 正在获取存储的输出(跳过连接)
-
x
=
self.act(l(x)) # 通过层和激活功能
-
-
return x
在 MNIST 上使用此架构训练的扩散模型生成以下示例(代码包含在补充材料中,但为简洁起见此处省略):
改进 UNet
这个简单的 UNet 可以完成这个相对简单的任务,但远非理想。那么,我们可以做些什么来改善它呢?
-
添加更多参数。这可以通过在每个块中使用多个卷积层、在每个卷积层中使用更多的过滤器或使网络更深来实现。
-
添加剩余连接。使用 ResBlocks 而不是常规的卷积层可以帮助模型学习更复杂的函数,同时保持训练稳定。
-
添加归一化,例如批量归一化。批量归一化可以帮助模型更快、更可靠地学习,方法是确保每一层的输出都以 0 为中心并且标准差为 1。
-
添加正则化,例如 dropout。Dropout 有助于防止模型过度拟合训练数据,这在处理较小的数据集时很重要。
-
加关注。通过引入自注意力层,我们允许模型在不同时间关注图像的不同部分,这可以帮助它学习更复杂的功能。添加类似 Transformer 的注意力层也让我们增加了可学习参数的数量,这可以帮助模型学习更复杂的功能。缺点是注意力层在高分辨率下比常规卷积层的计算成本高得多,因此我们通常只在较低分辨率下使用它们(即 UNet 中的较低分辨率块)。
-
为时间步添加一个额外的输入,以便模型可以根据噪声水平调整其预测。这称为时间步长条件,几乎用于所有最近的扩散模型。我们将在下一章看到更多关于条件模型的内容。
为了进行比较,以下是使用扩散器库中的 UNet 实现时 MNIST 的结果,它具有上述所有改进:
将来可能会使用结果和更多详细信息来扩展此部分。我们只是还没有开始训练具有不同改进的变体!
替代架构
最近,已经为扩散模型提出了许多替代架构。这些包括:
-
Transformers。Peebles 和 Xie 的DiT 论文(“Scalable Diffusion Models with Transformers”)表明,基于 Transformer 的架构可用于训练扩散模型,并取得了很好的效果。然而,变压器架构的计算和内存要求仍然是非常高分辨率的挑战。
-
来自Simple Diffusion 论文链接的UViT架构旨在通过用一大堆变压器块替换 UNet 的中间层来获得两全其美。本文的一个关键见解是,将大部分计算集中在 UNet 的较低分辨率块上,可以更有效地训练高分辨率扩散模型。对于非常高分辨率,他们使用称为小波变换的东西进行一些额外的预处理,以降低输入图像的空间分辨率,同时通过使用额外的通道保留尽可能多的信息,再次减少花费在更高的空间分辨率。
-
循环接口网络。RIN 论文(Jabri 等人)采用了类似的方法,首先将高分辨率输入映射到更易于管理和低维的“潜在”表示,然后由一堆变换器块处理,然后再解码回图像. 此外,RIN 论文引入了“循环”的概念,其中信息从先前的处理步骤传递到模型,这可能有利于扩散模型旨在执行的那种迭代改进。
基于转换器的方法是否会完全取代 UNet 作为扩散模型的首选架构,或者像 UViT 和 RIN 架构这样的混合方法是否会被证明是最有效的,还有待观察。
深入:目标和预处理
我们已经谈到了扩散模型采用噪声输入并“学习去噪”它。乍一看,您可能会假设网络的自然预测目标是图像的去噪版本,我们将其称为x0
。然而,在代码中,我们将模型预测与用于创建噪声版本的单位方差噪声(通常称为 epsilon 目标,eps
)进行了比较。这两者在数学上看起来是相同的,因为如果我们知道噪声和时间步长,我们就可以推导出来x0
,反之亦然。虽然这是事实,但目标的选择对不同时间步长的损失有多大有一些微妙的影响,从而影响模型学会最好地降噪的噪声水平。为了获得一些直觉,让我们想象一下跨不同时间步长的一些不同目标:
在极低的噪声水平下,x0
目标很容易实现,而准确预测噪声几乎是不可能的。同样,在极高的噪声水平下,eps
目标很容易,而准确预测去噪图像几乎是不可能的。这两种情况都不是理想的,因此引入了额外的目标,让模型预测不同时间步长的x0
混合eps
。目标(在Salimans 和 Ho 的“Progressive distillation for fast sampling of diffusion models”v
中介绍)就是这样一个目标,它被定义为v=√¯a·∊ + √1-¯a· x Karras 等人的EDM 论文通过称为 的参数引入了类似的想法c_skip
,并将不同的扩散模型公式统一到一个一致的框架中。如果您有兴趣了解更多有关不同扩散模型公式的不同目标、尺度和其他细微差别的信息,我们建议您阅读他们的论文以进行更深入的讨论。
项目时间:训练你自己的扩散模型
现在您已经了解了扩散模型的基础知识,是时候自己训练一些了!本章的补充材料包括一个笔记本,它会引导您完成在您自己的数据集上训练扩散模型的过程。在你完成它的过程中,回顾一下本章,看看不同的部分是如何组合在一起的。该笔记本还包含许多建议的更改,您可以进行更改以更好地探索不同的模型架构和训练策略如何影响结果。
概括
在本章中,我们了解了如何应用迭代优化的思想来训练能够将噪声转化为精美图像的扩散模型。您已经看到了一些用于创建成功扩散模型的设计选择,并希望通过训练您自己的模型将它们付诸实践。在下一章中,我们将了解一些更先进的技术,这些技术是为提高扩散模型的性能而开发的,并赋予它们非凡的新功能!