一文弄懂 Diffusion Model（DDPM）+ 代码实现

马鹏森

已于 2024-04-05 19:16:21 修改

阅读量9.5k

点赞数 37

分类专栏： Diffusion Model 文章标签：人工智能 python 算法

于 2022-11-27 09:17:53 首次发布

本文链接：https://blog.csdn.net/weixin_43135178/article/details/128060554

版权

Diffusion Model 专栏收录该内容

28 篇文章 71 订阅

订阅专栏

Denoising Diffusion Probabilistic Models（DDPM）

paper：https://arxiv.org/abs/2006.11239

整个扩散模型非常的简单，下面我所用的完整代码已经上传到“资源”了

★★★：十分钟读懂Diffusion：图解Diffusion扩散模型 - 知乎

一、简易版通过代码和有一定基础的直接阅读

网络训练目标

论文的源代码采用Unet实现噪声 eps 的预测，整个训练过程其实就是在训练Unet网络的参数

Unet职责

无论在前向过程还是反向过程，Unet的职责都是根据当前的样本和时间 t 预测噪声。

训练阶段是一步预测出从0到t加的噪声noise
推理阶段是根据模型算出该t时刻的噪声，然后用x_t减去该时刻的噪声noise得到x[t-1]

Algorithm1：Training

X0 --> Xt：

1）先看下总体如何进行训练的：

2）大致是如何训练的

3）详细步骤

从数据中抽取一个样本图片x0，
从1-T中随机选取一个时间t，并将时间步t 转化为对应embedding
将 x0 和 t embedding传给GaussionDiffusion，GaussionDiffusion采样一个随机噪声，加到 x0 ，形成xt ，然后将xt 和t放入Unet，Unet根据t正弦位置编码和 xt 结合，预测加的这个噪声，并返回预测噪声
将神经网络Unet预测的噪声与之前GaussionDiffusion采样的随机噪声求 loss，计算梯度，更新权重。
重复以上步骤，直到网络Unet训练完成。

代码实现参考自：https://github.com/yangqy1110/Diffusion-Models/blob/main/Diffusion%20Model.ipynb

1、定义数据

import matplotlib.pyplot as plt
import numpy as np
from sklearn.datasets import make_s_curve  # 生成S形二维数据点 https://scikit-learn.org/stable/modules/generated/sklearn.datasets.make_s_curve.html
import torch
import torch.nn as nn

## ----------------------------- 1、生成数据，(10000, 2)的数据点集，组成一个S形 ----------------------------- ##
s_curve, _ = make_s_curve(10 ** 4, noise=0.1)  # 10000个数据点
s_curve = s_curve[:, [0, 2]] / 10.0
print("shape of s:", np.shape(s_curve))
dataset = torch.Tensor(s_curve).float()

2、定义超参数的值

## ----------------------------- 2、确定超参数的值 ----------------------------- ##
# 采样时间步总长度 t
num_steps = 100

# 制定每一步的beta
betas = torch.linspace(-6, 6, num_steps)
betas = torch.sigmoid(betas) * (0.5e-2 - 1e-5) + 1e-5

# 计算alpha、alpha_prod、alpha_prod_previous、alpha_bar_sqrt等变量的值
alphas = 1 - betas
alphas_prod = torch.cumprod(alphas, 0)
alphas_prod_p = torch.cat([torch.tensor([1]).float(), alphas_prod[:-1]], 0)
alphas_bar_sqrt = torch.sqrt(alphas_prod)
one_minus_alphas_bar_log = torch.log(1 - alphas_prod)
one_minus_alphas_bar_sqrt = torch.sqrt(1 - alphas_prod)

3、确定扩散前向过程任意时刻的采样值 x[t]： x[0] + t --> x[t]

注意：这里每一timestep的加入的噪声是不一样的，是根据调度策略添加的噪声

## ----------------------------- 3、确定扩散前向过程任意时刻的采样值 x[t]： x[0] + t --> x[t] ----------------------------- ##
def q_x(x_0, t):
    """
    x[0] + t --> x[t]
    :param x_0:初始数据
    :param t:任意时刻
    :return:
    """
    noise = torch.randn_like(x_0)
    alphas_t = alphas_bar_sqrt[t]
    alphas_1_m_t = one_minus_alphas_bar_sqrt[t]
    x_t = alphas_t * x_0 + alphas_1_m_t * noise
    return x_t

4、编写预测噪声的U-Net

## ----------------------------- 4、编写求逆扩散过程噪声的模型U-Net（这里使用的是MLP模拟U-Net，官方使用的是U-Net） x[t] + t --> noise_predict----------------------------- ##
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):
        #  x = x[0]
        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

5、定义损失函数

## ----------------------------- 损失函数 = 真实噪声eps与预测出的噪声noise_predict 之间的loss ----------------------------- ##
def diffusion_loss_fn(model, x_0, alphas_bar_sqrt, one_minus_alphas_bar_sqrt, n_steps):
    """对任意时刻t进行采样计算loss"""
    batch_size = x_0.shape[0]

    # 对一个batchsize样本生成随机的时刻t, t的形状是torch.Size([batchsize, 1])
    t = torch.randint(0, n_steps, size=(batch_size // 2,))
    t = torch.cat([t, n_steps - 1 - t], dim=0)
    t = t.unsqueeze(-1)

    ## 1) 根据 alphas_bar_sqrt, one_minus_alphas_bar_sqrt --> 得到任意时刻t的采样值x[t]
    # x0的系数
    a = alphas_bar_sqrt[t]
    # 噪声eps的系数
    aml = one_minus_alphas_bar_sqrt[t]
    # 生成随机噪音eps
    e = torch.randn_like(x_0)
    # 得到任意时刻t的采样值
    x = x_0 * a + e * aml

    ## 2) x[t]送入U-Net模型，得到t时刻的随机噪声预测值
    output = model(x, t.squeeze(-1))

    ## 3)计算真实噪声eps与预测出的噪声之间的loss
    loss = (e - output).square().mean()
    return loss

6、训练模型

## ----------------------------- 训练模型 ----------------------------- ##
print('Training model...')
batch_size = 128
num_epoch = 4000
dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)

model = MLPDiffusion(num_steps)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for t in 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 % 100 == 0):
        print(loss)
        torch.save(model.state_dict(), 'model_{}.pth'.format(t))

Algorithm2：Sampling

Xt --> Xt-1：

1）整体感觉一下

2）详细步骤

从标准正态分布采样出与维度一致的 xT
从 T, T−1,...,2,1 依次重复以下步骤：
根据训练好的模型U-Net求出噪声，结合 xt 和采样得到z利用重参数化技巧，得到 xt−1
循环结束后返回 x0

cite : 超详细的扩散模型（Diffusion Models）原理+代码 - 知乎

1、定义单步逆扩散过程

功能：从x[t]采样t-1时刻的重构值x[t-1]

根据下面的算法公式得到x[t-1]

def p_sample(model, x, t, betas, one_minus_alphas_bar_sqrt):
    """
    从x[t]采样t-1时刻的重构值x[t-1]
    :param model:
    :param x: x[T]
    :param t:
    :param betas:
    :param one_minus_alphas_bar_sqrt:
    :return:
    """
    ## 1) 求出 bar_u_t
    t = torch.tensor([t])
    coeff = betas[t] / one_minus_alphas_bar_sqrt[t]
    # 送入U-Net模型，得到t时刻的随机噪声预测值 eps_theta
    eps_theta = model(x, t)
    mean = (1 / (1 - betas[t]).sqrt()) * (x - (coeff * eps_theta))

    ## 2) 得到 x[t-1]
    z = torch.randn_like(x)
    sigma_t = betas[t].sqrt()
    sample = mean + sigma_t * z
    return (sample)

2、根据单步逆扩散过程定义循环逆扩散过程

def p_sample_loop(model, noise_x_t, n_steps, betas, one_minus_alphas_bar_sqrt):
    """
    从x[T]恢复x[T-1]、x[T-2]|...x[0] 的循环
    :param model:
    :param shape:数据集的形状，也就是x[T]的形状
    :param n_steps:
    :param betas:
    :param one_minus_alphas_bar_sqrt:
    :return: x_seq由x[T]、x[T-1]、x[T-2]|...x[0]组成, cur_x是从噪声中生成的图片
    """
    # 得到噪声x[T]
    cur_x = noise_x_t
    x_seq = [noise_x_t]
    # 从x[T]恢复x[T-1]、x[T-2]|...x[0]
    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, cur_x

3、根据训练好的扩散模型 + 随机噪声 ==> 生成数据

# 1) 加载训练好的diffusion model
model = MLPDiffusion(num_steps)
model.load_state_dict(torch.load('model_1000.pth'))
# 2) 生成随机噪声x[T]
noise_x_t = torch.randn(dataset.shape)
# 3) 根据随机噪声逆扩散为x[T-1]、x[T-2]|...x[0] + 图片x[0]
x_seq, cur_x = p_sample_loop(model, noise_x_t, num_steps, betas, one_minus_alphas_bar_sqrt)
print(x_seq, cur_x)

二、带有公式的解析过程

一、前向 Diffusion 过程

Diffusion Model 首先定义了一个前向扩散过程，总共包含T个时间步，如下图所示：

最左边的蓝色圆圈 x0 表示真实自然图像，对应下方的狗子图片。

最右边的蓝色圆圈 xT 则表示纯高斯噪声，对应下方的噪声图片。

最中间的蓝色圆圈 xt 则表示加了噪声的 x0 ，对应下方加了噪声的狗子图片。

箭头下方的 q(xt|xt-1) 则表示一个以前一个状态 xt-1 为均值的高斯分布，xt 从这个高斯分布中采样得到。

所谓前向扩散过程可以理解为一个马尔可夫链[7]，即通过逐步对一张真实图片添加高斯噪声直到最终变成纯高斯噪声图片。

（1）利用前一时刻的 `xt-1` 得到任意时刻的噪声图片 `xt（`重参数化技巧`）`

那么具体是怎么添加噪声呢，公式表示如下：

★★★ 也就是每一时间步的 xt 是从一个，以 1-βt 开根号乘以 xt-1 为均值，βt为方差的高斯分布中采样得到的。其中βt, t ∈ [1, T] 是一系列固定的值，由一个公式生成。

在参考资料 [2] 中设置 T=1000, β1=0.0001, βT=0.02，并通过一句代码生成所有 βt 的值：

# https://pytorch.org/docs/stable/generated/torch.linspace.html
betas = torch.linspace(start=0.0001, end=0.02, steps=1000)

然后在采样得到 xt 的时候并不是直接通过高斯分布 q(xt|xt-1) 采样，而是用了一个重参数化的技巧（详见参考资料[4]第5页）。

★★★ 简单来说就是，如果想要从一个任意的均值 μ 方差 σ^2 的高斯分布中采样得到xt：

1）可以首先从一个标准高斯分布（均值0，方差1）中进行采样得到噪声 ε，

noise = torch.randn_like(x_0)

2）然后利用 μ + σ·ε 就等价于从任意均值 μ 方差 σ^2 的高斯分布中采样（首先从标准高斯分布中采样得到噪声 ε，接着乘以标准差再加上均值）。公式表示如下：

xt = sqrt(1-betas[t]) * xt-1 + sqrt(betas[t]) * noise

完整代码：

# https://pytorch.org/docs/stable/generated/torch.randn_like.html
betas = torch.linspace(start=0.0001, end=0.02, steps=1000)
noise = torch.randn_like(x_0)
xt = sqrt(1-betas[t]) * xt-1 + sqrt(betas[t]) * noise

（2）直接从 `x0` 采样得到中间任意一个时间步的噪声图片 `xt`

然后前向扩散过程还有个属性，就是可以直接从 x0 采样得到中间任意一个时间步的噪声图片 xt，公式如下：

其中的 αt 表示：

具体怎么推导出来的可以看参考资料[4] 第11页，伪代码表示如下：

betas = torch.linspace(start=0.0001, end=0.02, steps=1000)
alphas = 1 - betas
# cumprod 相当于为每个时间步 t 计算一个数组 alphas 的前缀乘结果
# https://pytorch.org/docs/stable/generated/torch.cumprod.html
alphas_cum = torch.cumprod(alphas, 0)
alphas_cum_s = torch.sqrt(alphas_cum)
alphas_cum_sm = torch.sqrt(1 - alphas_cum)

# 应用重参数化技巧采样得到 xt
noise = torch.randn_like(x_0)
xt = alphas_cum_s[t] * x_0 + alphas_cum_sm[t] * noise

通过上述的讲解，读者应该对 Diffusion Model 的前向扩散过程有比较清晰的理解了。

不过我们的目的不是要做图像生成吗？现在只是从数据集中的真实图片得到一张噪声图，那具体是怎么做图像生成呢？

二、反向 Diffusion 过程

反向扩散过程 q(xt-1|xt, x0) （看粉色箭头）是前向扩散过程 q(xt|xt-1) 的后验概率分布。

和前向过程相反是从最右边的纯高斯噪声图，逐步采样得到真实图像 x0。

后验概率 q(xt-1|xt, x0) 的形式可以根据贝叶斯公式推导得到（推导过程详见参考资料[4]第12页）：

也是一个高斯分布。

（1）方差：

其方差从公式上看是个常量，所有时间步的方差值都是可以提前计算得到的：

betas = torch.linspace(start=0.0001, end=0.02, steps=1000)
alphas = 1 - betas
alphas_cum = torch.cumprod(alphas, 0)
alphas_cum_prev = torch.cat((torch.tensor([1.0]), alphas_cum[:-1]), 0)
posterior_variance = betas * (1 - alphas_cum_prev) / (1 - alphas_cum)

（2）均值：

对于反向扩散过程，在采样生成 xt-1 的时候 xt 是已知的，而其他系数都是可以提前计算得到的常量。

但是现在问题来了，在真正通过反向过程生成图像的时候，x0 我们是不知道的，因为这是待生成的目标图像。

好像变成了鸡生蛋，蛋生鸡的问题，那该怎么办呢？

（3）Diffusion Model 训练目标

当一个概率分布q 求解困难的时候，我们可以换个思路（详见参考资料[5,6]）。

通过人为构造一个新的分布 p，然后目标就转为缩小分布 p 和 q 之间差距。通过不断修改 p 的参数去缩小差距，当 p 和 q 足够相似的时候就可以替代 q 了。

然后回到反向 Diffusion 过程，由于后验分布 q(xt-1|xt, x0) 没法直接求解。

那么我们就构造一个高斯分布 p(xt-1|xt)（见绿色箭头），让其方差和后验分布 q(xt-1|xt, x0) 一致：

而其均值则设为：

和 q(xt-1|xt, x0) 的区别在于，x0 改为 xθ(xt, t) 由一个深度学习模型预测得到，模型输入是噪声图像 xt 和时间步 t 。

然后缩小分布 p(xt-1|xt) 和 q(xt-1|xt, x0) 之间差距，变成优化以下目标函数（推导过程详见参考资料[4]第13页）：

但是如果让模型直接从 xt 去预测 x0，这个拟合难度太高了，我们再继续换个思路。

前面介绍前向扩散过程提到，xt 可以直接从 x0 得到：

将上面的公式变换一下形式：

代入上面 q(xt-1|xt, x0) 的均值式子中可得（推导过程详见参考资料[4]第15页）：

观察上述变换后的式子，发现后验概率 q(xt-1|xt, x0) 的均值只和 xt 和前向扩散时候时间步 t 所加的噪声有关。

所以我们同样对构造的分布 p(xt-1|xt) 的均值做一下修改：

将模型改为去预测在前向时间步 t 所添加的高斯噪声 ε，模型输入是 xt 和时间步 t：

接着优化的目标函数就变为（推导过程详见参考资料[4]第15页）：

然后训练过程算法描述如下，最终的目标函数前面的系数都去掉了，因为是常量：

★ 可以看到虽然前面的推导过程很复杂，但是训练过程却很简单：

首先每个迭代就是从数据集中取真实图像 x0，并从均匀分布中采样一个时间步 t，
然后从标准高斯分布中采样得到噪声 ε，并根据公式计算得到前向过程的 xt。
接着将 xt 和 t 输入到模型让其输出去拟合预测噪声 ε，并通过梯度下降更新模型，一直循环直到模型收敛。
而采用的深度学习模型是类似 UNet 的结构（详见参考资料[2]附录B）。

训练过程的伪代码如下：

betas = torch.linspace(start=0.0001, end=0.02, steps=1000)
alphas = 1 - betas
alphas_cum = torch.cumprod(alphas, 0)
alphas_cum_s = torch.sqrt(alphas_cum)
alphas_cum_sm = torch.sqrt(1 - alphas_cum)

def diffusion_loss(model, x0, t, noise):
    # 根据公式计算 xt
    xt = alphas_cum_s[t] * x0 + alphas_cum_sm[t] * noise
    # 模型预测噪声
    predicted_noise = model(xt, t)
    # 计算Loss
    return mse_loss(predicted_noise, noise)

for i in len(data_loader):
    # 从数据集读取一个 batch 的真实图片
    x0 = next(data_loader)
    # 采样时间步
    t = torch.randint(0, 1000, (batch_size,))
    # 生成高斯噪声
    noise = torch.randn_like(x_0)
    loss = diffusion_loss(model, x0, t, noise)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

三、Diffusion Model 生成图像过程（推理）

模型训练好之后，在真实的推理阶段就必须从时间步 T 开始往前逐步生成图片，算法描述如下：

步骤：

一开始先生成一个从标准高斯分布生成噪声，
然后每个时间步 t，将上一步生成的图片 xt 输入模型模型预测出噪声。
接着从标准高斯分布中采样一个噪声，根据重参数化技巧，后验概率的均值和方差公式，计算得到 xt-1，直到时间步 1 为止。

参考资料

[1] https://www.assemblyai.com/blog/diffusion-models-for-machine-learning-introduction/
[2] https://arxiv.org/pdf/2006.11239.pdf
[3] https://arxiv.org/pdf/2102.09672.pdf
[4] https://arxiv.org/pdf/2208.11970.pdf
[5] https://www.zhihu.com/question/41765860/answer/1149453776
[6] https://www.zhihu.com/question/41765860/answer/331070683
[7] https://zh.wikipedia.org/wiki/%E9%A9%AC%E5%B0%94%E5%8F%AF%E5%A4%AB%E9%93%BE
[8] https://github.com/rosinality/denoising-diffusion-pytorch
[9] https://github.com/openai/improved-diffusion