结合案例白话 diffusion 模型原理-六虎

本文正在参与「金石计划」

前语

本文的首要内容是运用文言介绍分散模型 diffusion 原理，而且一起结合代码简单演示 diffusion 模型的功用，咱们运用一个随机生成的“S”形图画，建立一个“简易版的 U-Net 模型”，先加噪练习模型，然后运用练习好的模型再去噪的进程。

数据准备

这儿是直接运用 make_s_curve 来生成一个包含了 10000 个点的“S”形图画，服从标准差为 0.1 的高斯噪声点集合，由于每个点的数据都是一个三维的，咱们从里边抽取第 0 维和第 2 维，形成一个二维的“S”图片，如下所示，用这个作为咱们的数据集。

import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import numpy as np
from sklearn.datasets import make_s_curve
import io
from PIL import Image
from tqdm import tqdm
s_curve,_ = make_s_curve(10**4, noise=0.1)
s_curve = s_curve[:, [0,2]]/10.0
dataset = torch.Tensor(s_curve).float()
print("数据集的巨细为:", np.shape(dataset))
# 绘制生成的 “S” 形图画
data = s_curve.T
fig,ax = plt.subplots()
ax.scatter(*data, color='blue', edgecolor='white');
ax.axis('off')
plt.show()

打印成果：

数据集的巨细为: torch.Size([10000, 2])

前向进程

这儿展现的是前向进程，公式如下，咱们能够直接运用初始图画来得到某个时间步的带噪音图画，也便是下面的 q_x 函数：

q(xt∣x0)=N(xt;tx0,(1−t)I)q(\mathbf{x}_{t}\mid\mathbf{x}_{0}) = \mathcal{N}(\mathbf{x}_{t} ; \sqrt{\bar{\alpha}_{t}}\mathbf{x}_{0},(1-\bar{\alpha}_{t})\mathbf{I})

其实写成人话便是，这些变量都是之前已经核算好的，能够直接拿来用， $x0\mathbf{x}_{0}$ 表明的便是开端的图片， $I\mathbf{I}$ 便是由高斯散布发生的噪声：

q(xt∣x0)=tx0+(1−t)Iq(\mathbf{x}_{t}\mid\mathbf{x}_{0}) = \sqrt{\bar{\alpha}_{t}}\mathbf{x}_{0}+\sqrt{(1-\bar{\alpha}_{t})}\mathbf{I}

咱们定义一共前向分散进程有 num_steps 步，运用 torch.linspace 来初始化 num_steps 个 $\beta$ 值，取值范围区间在 start 和 end 上均匀距离的 num_steps 个数，这样咱们就得到了每个 step 时分的 $\beta$ 值，依据如下公式，间接地咱们也一起得到了每个 step 时分的 $\alpha$ 值。

t=1−t\alpha_{t} = 1-\beta_{t}

由于后续的在前向分散进程中在生成某个 step 时分的噪声图是能够直接用开端的图画样本核算得到的，而这些核算需求的 $t\bar{\alpha}_{t}$ 、 $t\sqrt{\bar{\alpha}_{t}}$ 、 $(1−t)\sqrt{(1-\bar{\alpha}_{t})}$ 等，所以能够提早核算出来。

这儿咱们展现了在原始图画上加噪声的前 10 个 step 图画成果，能够看出来在到了第 5 个 step 之后基本就含糊的连他 mother 都认不出来了。

def make_beta_schedule(n_timesteps=1000, start=1e-5, end=1e-2):
    betas = torch.linspace(start, end, n_timesteps)
    return betas
# 提早核算需求的各种变量
num_steps = 100
betas = make_beta_schedule(n_timesteps=num_steps, start=1e-5, end=0.5e-2)
alphas = 1-betas
alphas_prod = torch.cumprod(alphas, 0)
alphas_bar_sqrt = torch.sqrt(alphas_prod)
one_minus_alphas_bar_sqrt = torch.sqrt(1 - alphas_prod)
# 前向加噪声进程
def q_x(x_0, t):
    noise = torch.randn_like(x_0)
    alphas_t = alphas_bar_sqrt[t]
    alphas_1_m_t = one_minus_alphas_bar_sqrt[t]
    return (alphas_t * x_0 + alphas_1_m_t * noise)
# 展现将图片加噪声的前 10 个成果
num_shows = 10
fig,axs = plt.subplots(2, 5, figsize=(10,5))
plt.rc('text',color='black')
for i in range(num_shows):
    j = i//5
    k = i%5
    q_i = q_x(dataset, torch.tensor([i]))  
    axs[j,k].scatter(q_i[:,0],q_i[:,1],color='red',edgecolor='white')
    axs[j,k].set_axis_off()
    axs[j,k].set_title('$q(\mathbf{x}_{'+str(i+1)+'})$')

这儿定义了前向仅仅用最简单的方法构建了一个相似 U-Net 的模型，运用的练习集便是上面带噪声的图片，标签便是方才咱们人为加的噪声，需求猜测的便是这些噪声散布，尽量减小实在噪声和猜测噪声的差错。需求留意的是在进行前向传达的时分，也加入了位置编码信息。用来学习不同 step 时分的噪声。

class MLPDiffusion(nn.Module):
    def __init__(self, n_steps, num_units=128):
        super(MLPDiffusion, self).__init__()
        self.linears = nn.ModuleList(
            [
                nn.Linear(2, num_units),
                nn.ReLU(),
                nn.Linear(num_units, num_units),
                nn.ReLU(),
                nn.Linear(num_units, num_units),
                nn.ReLU(),
                nn.Linear(num_units, 2),
            ]
        )
        self.step_embeddings = nn.ModuleList(
            [
                nn.Embedding(n_steps, num_units),
                nn.Embedding(n_steps, num_units),
                nn.Embedding(n_steps, num_units),
            ]
        )
    def forward(self, x, t):
        for idx, embedding_layer in enumerate(self.step_embeddings):
            t_embedding = embedding_layer(t)
            x = self.linears[2 * idx](x)
            x += t_embedding
            x = self.linears[2 * idx + 1](x)
        x = self.linears[-1](x)
        return x

这儿是核算丢失值的丢失函数，其实便是依据上面相同的公式来核算某个 step 时分的带噪图片 x ，将这个带噪图片 x 输入模型中，让模型猜测出噪声 output ，然后核算 x 和 output 两者的均方差作为丢失值，咱们只需求练习模型不断削减这个丢失即可。

q(xt∣x0)=tx0+(1−t)Iq(\mathbf{x}_{t}\mid\mathbf{x}_{0}) = \sqrt{\bar{\alpha}_{t}}\mathbf{x}_{0}+\sqrt{(1-\bar{\alpha}_{t})}\mathbf{I}

def diffusion_loss_fn(model, x_0, alphas_bar_sqrt, one_minus_alphas_bar_sqrt, n_steps):
    batch_size = x_0.shape[0]
    t = torch.randint(0, n_steps, size=(batch_size // 2,))
    t = torch.cat([t, n_steps - 1 - t], dim=0)
    t = t.unsqueeze(-1)
    a = alphas_bar_sqrt[t]
    aml = one_minus_alphas_bar_sqrt[t]
    e = torch.randn_like(x_0)
    x = x_0 * a + e * aml
    output = model(x, t.squeeze(-1))
    return (e - output).square().mean()

反向进程

这儿定义了反向复原进程，咱们将带噪声的图片从最终的含糊模样，不断进行迭代去噪的进程，最终复原回咱们的原始图片。这儿需求留意的是依照原论文中规则了某个 step 的方差 $\sigma^{2}_{t}$ = $\beta_{t}$ ，所以咱们能够直接进行方差的核算。

这个进程应该是会比较慢的，由于要从最终一步开端，一步一步地向前进行每一步的去噪操作。p_sample_loop 最终的返回成果巨细是 [101, 10000, 2] ，只要第一个元素是待去噪的原始含糊图画输入，剩下的 100 个都是反向进程中每个 step 去噪之后的图画成果。

def p_sample(model, x, t, betas, one_minus_alphas_bar_sqrt):
    t = torch.tensor([t])
    coeff = betas[t] / one_minus_alphas_bar_sqrt[t]
    eps_theta = model(x, t)
    mean = (1 / (1 - betas[t]).sqrt()) * (x - (coeff * eps_theta))
    z = torch.randn_like(x)
    sigma_t = betas[t].sqrt()
    sample = mean + sigma_t * z
    return (sample)
def p_sample_loop(model, shape, n_steps, betas, one_minus_alphas_bar_sqrt):
    cur_x = torch.randn(shape)
    x_seq = [cur_x]
    for i in reversed(range(n_steps)):
        cur_x = p_sample(model, cur_x, i, betas, one_minus_alphas_bar_sqrt)
        x_seq.append(cur_x)
    return x_seq

练习

这儿是开端练习模型的进程，为了能体现出作用，需求阅历 5000 个 epoch ，每个 batch 为 128 ，运用学习率为 5*1e-4 的 Adam 作为咱们的优化器。

遍历每一个 batch ，需求留意的是，每个 batch 中的样本是从 dataset 中随机抽取的 batch_size 个二维点位信息组成的图画，也便是每个样本不是一开端的 “S” 形状，而是各种由 128 点组成的图画，由于这 128 个点都是从原始“S”图中采样得到的，所以也基本都是“S”形状的（如下展现出一个样本图片），这首要是为了简单生成丰富咱们的数据集。先经过模型前向传达核算丢失，然后反向传达更新模型权重参数。

plt.rc('text',)
print("dataset 的 shape 是：",dataset.shape)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)
for idx, batch_x in enumerate(dataloader):
    print("batch 的 shape 是:",batch_x.shape)
    fig,ax = plt.subplots()
    ax.scatter(*batch_x.T, color='blue', edgecolor='white');
    ax.axis('off')
    plt.show()
    break

成果打印：

dataset 的 shape 是： torch.Size([10000, 2])
batch 的每个样本的 shape 是: torch.Size([128, 2])

每阅历若干个 epoch ，打印一个丢失值，而且运用猜测好的模型进行一次反向进程，将反向进程中发生的去噪成果图片每隔 10 个 step 展现出来，从左到右的方向展现了从带噪图片输入到不断去噪进程。一共有 5000 个 epoch ，所以显示出来 5 行成果，咱们能够看到在第 5 行已经初见雏形，如果 epoch 再增大，作用会愈加显着。

plt.rc('text', color='blue')
batch_size = 128
num_epoch = 5000
print("dataset 的 shape 是：",dataset.shape)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)
model = MLPDiffusion(num_steps)   
optimizer = torch.optim.Adam(model.parameters(), lr=5*1e-4)
for t in tqdm(range(num_epoch)):
    for idx, batch_x in enumerate(dataloader):
        loss = diffusion_loss_fn(model, batch_x, alphas_bar_sqrt, one_minus_alphas_bar_sqrt, num_steps)
        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.)
        optimizer.step()
    if (t % 1000 == 0):
        print("第 %d 个 epoch 的丢失值为 %f "%(t,loss))
        x_seq = p_sample_loop(model, dataset.shape, num_steps, betas, one_minus_alphas_bar_sqrt)  # [101, 10000, 2]
        fig, axs = plt.subplots(1, 10, figsize=(28, 3))
        for i in range(1, 11):
            cur_x = x_seq[i * 10].detach()
            axs[i - 1].scatter(cur_x[:, 0], cur_x[:, 1], color='red', edgecolor='white');
            axs[i - 1].set_axis_off();
            axs[i - 1].set_title('$q(\mathbf{x}_{' + str(i * 10) + '})$')

成果打印：

dataset 的 shape 是： torch.Size([10000, 2])

这儿咱们分别将前向进程的前 10 个 step 的加噪进程图画、和反向进程的 50 个去噪进程做成 GIF 愈加形象地展现出 diffusion 模型的两个进程作用。

imgs = []
for i in tqdm(range(10)):
    plt.clf()
    q_i = q_x(dataset, torch.tensor([i]))
    plt.scatter(q_i[:, 0], q_i[:, 1], color='red', edgecolor='white', s=5);
    plt.axis('off');
    img_buf = io.BytesIO()
    plt.savefig(img_buf, format='png')
    img = Image.open(img_buf)
    imgs.append(img)
imgs[0].save("前向.gif",format='GIF', append_images=imgs, save_all=True, duration=300, loop=0)

reverse = []
for i in tqdm(range(0,100,2)):
    plt.clf()
    cur_x = x_seq[i].detach()
    plt.scatter(cur_x[:, 0], cur_x[:, 1], color='red', edgecolor='white', s=5);
    plt.axis('off')
    img_buf = io.BytesIO()
    plt.savefig(img_buf, format='png')
    img = Image.open(img_buf)
    reverse.append(img)
reverse[0].save("反向.gif",format='GIF',append_images=reverse, save_all=True ,duration=200, loop=0)

到这儿就介绍完了，不足之处请多见谅，欢迎相互讨论学习。

参阅

github.com/cat-meowmeo…
huggingface.co/blog/annota… zhuanlan.zhihu.com/p/572161541
zhuanlan.zhihu.com/p/572161541

结合案例白话 diffusion 模型原理

前语

相关常识简介

数据准备

前向进程

反向进程

练习

参阅

作者信息

结合案例白话 diffusion 模型原理

前语

相关常识简介

数据准备

前向进程

反向进程

练习

参阅

相关文章

初识 Swift：字符串示例

前端必备的定时任务技能 – Cron + node-schedule

为什么 Rust 中的枚举(enum)很常见很好用? ? ?

Golang 内存调优 – 逃逸分析

作者信息