【笔记】扩散模型（一）：DDPM 理论推导与代码实现

论文链接：Denoising Diffusion Probabilistic Models

在正式开始之前先贴一个生成结果示例：

端午假期卷一卷，开一个新坑系统性地整理一下扩散模型的相关知识。扩散模型的名字来源于物理中的扩散过程，对于一张图像来说，类比扩散过程，向这张图像逐渐加入高斯噪声，当加入的次数足够多的时候，图像中像素的分布也会逐渐变成一个高斯分布。当然这个过程也可以反过来，如果我们设计一个神经网络，每次能够从图像中去掉一个高斯噪声，那么最后就能从一个高斯噪声得到一张图像。虽然一张有意义的图像不容易获得，但高斯噪声很容易采样，如果能实现这个逆过程，就能实现图像的生成。

这个过程可以形象地用上图表示，扩散模型中有两个过程，分别是前向过程（从图像加噪得到噪音）和反向过程（从噪音去噪得到图像）。在上图中，向图像 ${\mathbf{x}}_0$ 逐渐添加噪声可以得到一系列的 ${\mathbf{x}}_1,{\mathbf{x}}_2,...,{\mathbf{x}}_T$，最后的 ${\mathbf{x}}_T$ 即接近完全的高斯噪声，这个过程显然是比较容易的。而从 ${\mathbf{x}}_T$ 逐渐去噪得到 ${\mathbf{x}}_0$​ 并不容易，扩散模型学习的就是这个去噪的过程。

前向过程

我们从比较简单的前向过程开始，第一个问题是如何向图像中添加高斯噪声。在 DDPM 中，加噪的方式是直接对图像和标准高斯噪声 ${ϵ}_{t-1}\sim \mathcal{N}(0,\mathbf{I})$ 进行加权求和：
$${\mathbf{x}}_t=\sqrt{1-{\beta }_t}{\mathbf{x}}_{t-1}+\sqrt{{\beta }_t}{ϵ}_{t-1}$$
这里的 ${\beta }_t$ 就是每一步加噪使用的方差，在实际上进行加噪时，起始时使用的方差比较小，随着加噪步骤增加，方差会逐渐增大。例如在 DDPM 的原文中，使用的方差是从 ${\beta }_1=10^{-4}$ 随加噪时间步线性增大到 ${\beta }_T=0.02$。这样设置主要是为了方便模型进行学习，如果在最开始就加入很大的噪声，对图像信息的破坏会比较严重，不利于模型学习图像的信息。这个过程也可以从反向进行理解，即去噪时先去掉比较大的噪音得到图像的雏形，再去掉小噪音进行细节的微调。

在上边的公式里，我们可以认为 ${\mathbf{x}}_t$ 满足均值为 $\sqrt{1-{\beta }_t}{\mathbf{x}}_{t-1}$，方差为 $\sqrt{{\beta }_t}\mathbf{I}$ 的高斯分布。这样可以把上述加权求和的过程写成条件概率分布的形式：
$$q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})=\mathcal{N}({\mathbf{x}}_t;\sqrt{1-{\beta }_t}{\mathbf{x}}_{t-1},{\beta }_t\mathbf{I})$$
上边等号的右边表示的就是当前的变量 ${\mathbf{x}}_t$ 满足一个 $\mathcal{N}(\sqrt{1-{\beta }_t}{\mathbf{x}}_{t-1},{\beta }_t\mathbf{I})$ 的概率分布。通过上边的公式我们可以看到，每一个时间步的 ${\mathbf{x}}_t$ 都只和 ${\mathbf{x}}_{t-1}$ 有关，因此这个扩散过程是一个马尔可夫过程。在前向过程中，每一步的 $\beta$ 都是固定的，真正的变量只有 ${\mathbf{x}}_{t-1}$，那么我们可以将公式中的 ${\mathbf{x}}_{t-1}$ 进一步展开：
$$\begin{array}{rl}{\mathbf{x}}_t& =\sqrt{1-{\beta }_t}{\mathbf{x}}_{t-1}+\sqrt{{\beta }_t}{ϵ}_{t-1}\\ & =\sqrt{1-{\beta }_t}(\sqrt{1-{\beta }_{t-1}}{\mathbf{x}}_{t-2}+\sqrt{{\beta }_{t-1}}{ϵ}_{t-2})+\sqrt{{\beta }_t}{ϵ}_{t-1}\\ & =\sqrt{(1-{\beta }_t)(1-{\beta }_{t-1})}{\mathbf{x}}_{t-2}+\sqrt{(1-{\beta }_t){\beta }_{t-1}}{ϵ}_{t-2}+\sqrt{{\beta }_t}{ϵ}_{t-1}\end{array}$$
在上边的公式里，实际上 ${ϵ}_{t-2}$ 和 ${ϵ}_{t-1}$ 是同分布的，都是 $\mathcal{N}(0,1)$，因此可以进行合并：
$$\begin{array}{rl}{\mathbf{x}}_t& =\sqrt{(1-{\beta }_t)(1-{\beta }_{t-1})}{\mathbf{x}}_{t-2}+\sqrt{(\sqrt{(1-{\beta }_t){\beta }_{t-1}}{)}^2+(\sqrt{{\beta }_t}{)}^2}{\overline{ϵ}}_{t-2}\\ & =\sqrt{(1-{\beta }_t)(1-{\beta }_{t-1})}{\mathbf{x}}_{t-2}+\sqrt{1-(1-{\beta }_t)(1-{\beta }_{t-1})}{\overline{ϵ}}_{t-2}\end{array}$$
令 ${\alpha }_t=1-{\beta }_t$，${\bar{\alpha }}_t={\prod }_{i=1}^t{\alpha }_i$，继续推导，可以得到：
$$\begin{array}{rl}{\mathbf{x}}_t& =\sqrt{{\alpha }_t{\alpha }_{t-1}}{\mathbf{x}}_{t-2}+\sqrt{1-{\alpha }_t{\alpha }_{t-1}}{\overline{ϵ}}_{t-2}\\ & =\cdots \\ & =\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ\end{array}$$
通过上述的推导，我们发现给定 ${\mathbf{x}}_0$ 和加噪的时间步，可以直接用一步就得到 ${\mathbf{x}}_t$，而并不需要一步步地重复最开始的加权求和。和上述同理，这个关系也可以写成：
$$q({\mathbf{x}}_t|{\mathbf{x}}_0)=\mathcal{N}({\mathbf{x}}_t;\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0,(1-{\bar{\alpha }}_t)\mathbf{I})$$
从这个式子里我们可以看出，加噪过程中的 ${\mathbf{x}}_t$ 可以看作原始图像 ${\mathbf{x}}_0$ 和高斯噪声 $ϵ$ 的线性组合，且两个组合系数的平方和为 1。在实现加噪过程时，加噪的 scheduler 也是根据 ${\bar{\alpha }}_t$ 设计的，这样更加直接，且为了保证最后得到的足够接近噪声，可以将 ${\bar{\alpha }}_t$ 直接设置为一个接近 0 的数。

反向过程

正如文章开始所说的，反向过程就是从 ${\mathbf{x}}_T$ 逐渐去噪得到 ${\mathbf{x}}_0$ 的过程，也就是求 $q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$。根据贝叶斯公式：
$$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)=\frac{q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})q({\mathbf{x}}_{t-1})}{q({\mathbf{x}}_t)}$$
在上边的公式里，在前文中我们已经给出了 $q(\mathbf{x}|{\mathbf{x}}_{t-1})$，但 $q({\mathbf{x}}_{t-1})$ 和 $q({\mathbf{x}}_t)$ 依然是未知的。虽然这两个分布目前未知，但是在上一节的最后，我们已经推导出了 $q({\mathbf{x}}_t|{\mathbf{x}}_0)$ 这个分布，那么我们可以给上面的贝叶斯公式加上 ${\mathbf{x}}_0$ 作为条件，将等号右侧的两个未知分布转化为已知分布：
$$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)=\frac{q({\mathbf{x}}_t|{\mathbf{x}}_{t-1},{\mathbf{x}}_0)q({\mathbf{x}}_{t-1}|{\mathbf{x}}_0)}{q({\mathbf{x}}_t|{\mathbf{x}}_0)}$$
而且因为先验分布 $q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})$ 是马尔可夫过程，${\mathbf{x}}_t$ 只与 ${\mathbf{x}}_{t-1}$ 有关，而与 ${\mathbf{x}}_0$ 无关，所以上边式子里的 $q({\mathbf{x}}_t|{\mathbf{x}}_{t-1},{\mathbf{x}}_0)=q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})$。但推导到这里还有问题，我们把 ${\mathbf{x}}_0$ 加入到了条件概率分布的条件中，但 ${\mathbf{x}}_0$ 依然是未知的，因此我们需要继续推导出一个与 ${\mathbf{x}}_0$ 无关的式子。

上面的公式右侧的几个条件概率分布全都是高斯分布：
$$\begin{array}{rl}q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})& =\mathcal{N}({\mathbf{x}}_t;\sqrt{{\alpha }_t}{\mathbf{x}}_{t-1},1-{\alpha }_t)\\ q({\mathbf{x}}_{t-1}|{\mathbf{x}}_0)& =\mathcal{N}({\mathbf{x}}_{t-1};\sqrt{{\bar{\alpha }}_{t-1}}{\mathbf{x}}_0,1-{\bar{\alpha }}_{t-1})\\ q({\mathbf{x}}_t|{\mathbf{x}}_0)& =\mathcal{N}({\mathbf{x}}_t;\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0,1-{\bar{\alpha }}_t)\end{array}$$
用概率密度函数把这个公式展开，如果不看前边的常数项，可以得到：
$$\begin{array}{rl}q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)& \propto \exp \left(-\frac{1}{2}\left[\frac{({\mathbf{x}}_t-\sqrt{{\alpha }_t}{\mathbf{x}}_{t-1}{)}^2}{{\beta }_t}+\frac{({\mathbf{x}}_{t-1}-\sqrt{{\bar{\alpha }}_{t-1}}{\mathbf{x}}_0{)}^2}{1-{\bar{\alpha }}_{t-1}}+\frac{({\mathbf{x}}_t-\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0)}{1-{\bar{\alpha }}_t}\right]\right)\end{array}$$
因为我们在这一步去噪的时候想求得的是 ${\mathbf{x}}_{t-1}$ 的分布，所以我们把上式展开并整理成一个关于 ${\mathbf{x}}_{t-1}$ 的多项式：
$$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)\propto \exp \left(-\frac{1}{2}\left[\left(\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\bar{\alpha }}_{t-1}}\right){\mathbf{x}}_{t-1}^2-\left(\frac{2\sqrt{{\alpha }_t}}{{\beta }_t}{\mathbf{x}}_t+\frac{2\sqrt{{\bar{\alpha }}_{t-1}}}{1-{\bar{\alpha }}_{t-1}}{\mathbf{x}}_0\right){\mathbf{x}}_{t-1}+C({\mathbf{x}}_t,{\mathbf{x}}_0)\right]\right)$$
上边的式子里常数项不重要（因为可以直接变成常数从指数部分挪走），所以可以暂时不管。对比高斯分布（可以证明反向过程的分布也是高斯分布）的指数部分 $\exp \left(-\frac{1}{2}\left(\frac{1}{{\sigma }^2}{x}^2-\frac{2\mu }{{\sigma }^2}x+\frac{{\mu }^2}{{\sigma }^2}\right)\right)$：
$$\{\begin{array}{l}\begin{array}{rl}\frac{1}{{\sigma }^2}& =\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\bar{\alpha }}_{t-1}}\\ \frac{2\mu }{{\sigma }^2}& =\frac{2\sqrt{{\alpha }_t}}{{\beta }_t}{\mathbf{x}}_t+\frac{2\sqrt{{\bar{\alpha }}_{t-1}}}{1-{\bar{\alpha }}_{t-1}}{\mathbf{x}}_0\end{array}\end{array}$$
可以发现 $\sigma$ 的表达式里都是我们 scheduler 里的定值，而求解出均值 $\mu$：
$$\mu =\frac{\sqrt{{\alpha }_t}(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}{\mathbf{x}}_t+\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t}{1-{\bar{\alpha }}_t}{\mathbf{x}}_0$$
代入上一章最后的 ${\mathbf{x}}_t=\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ$，得到：
$$\mu =\frac{1}{\sqrt{{\alpha }_t}}\left({\mathbf{x}}_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\widetilde{ϵ}}_t\right)$$
注意在反向过程中我们并不知道在前向过程中加入的噪声 ${ϵ}_t$ 是 $\mathcal{N}(0,1)$ 中的具体哪一个噪声，而噪声也没有办法继续转换成其他的形式。因此我们使用神经网络在反向过程中估计的目标就是 ${\widetilde{ϵ}}_t$。在这个网络中，输入除了 ${\mathbf{x}}_t$ 之外还需要 $t$，可以简单理解为：加噪过程中 ${\mathbf{x}}_t$ 的噪声含量是由 $t$ 决定的，因此在预测噪声时也需要知道时间步 $t$​ 作为参考，以降低预测噪声的难度。

注：关于反向过程为什么要这样做，Lilian Weng 基于变分推断给出了一个复杂的证明，因为过于难以理解，这里暂且把它跳过。（以后有可能会填坑，也有可能不会x）

具体的训练过程

我们已经知道了去噪网络的参数和预测目标，下一个问题就是如何去训练这个去噪网络。原始论文中给出了如下的训练过程：

在上面的算法中，首先从数据集 $q({\mathbf{x}}_0)$ 中采样出 ${\mathbf{x}}_0$，从 1 到 T 的均匀分布中采样出 $t$，从标准高斯分布中采样出 $ϵ$。然后根据 ${\mathbf{x}}_t=\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ$ 将 ${\mathbf{x}}_0$ 与 $ϵ$ 加权求和得到噪声图，最后将噪声图和时间步输入到网络中预测噪声，并用真实的噪声计算出 L2 损失进行优化。

这里比较难理解的是 $ϵ$ 本身就是从标准高斯分布中采样出的，为什么还需要一个网络专门对其进行预测。我个人的理解是：尽管每次添加的噪声都是从固定的分布中采样出的，但如果用同一个分布中的另一个采样出的样本将其代替，就会向去噪过程引入一定的误差，最后这些误差积累的结果会破坏最终生成的图像。

具体的采样过程

论文中同样也给出了采样过程：

具体来说，首先从标准正态分布中采样出 ${\mathbf{x}}_T$ 作为初始的图像，然后重复 $T$ 步去噪过程。在每一步去噪过程中，由于我们已经推导出：
$$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)=\mathcal{N}\left({\mathbf{x}}_{t-1};\frac{1}{\sqrt{{\alpha }_t}}\left({\mathbf{x}}_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\widetilde{ϵ}}_t\right),{\sigma }_t^2\right)$$
利用一个重参数化技巧：从 $\mathcal{N}(\mu ,{\sigma }^2)$ 采样可以实现为从 $\mathcal{N}(0,1)$ 采样出 $ϵ$，再计算 $\mu +ϵ\cdot \sigma$。这样即可实现从上述的高斯分布中采样出 ${\mathbf{x}}_{t-1}$。如此重复 $T$ 次即可得到最终的结果，注意最后一步的时候没有采样，而是只加上了均值。

DDPM 的代码实现

现有的主流方法使用 UNet 来实现去噪网络，如下图所示。

为了降低理解的难度，我们这里不关心这个去噪网络的具体实现，只需要知道这个网络接收一个噪声图 ${\mathbf{x}}_t$ 和一个时间步 $t$ 作为参数，并输出一个噪声的预测结果 ${ϵ}_{\theta }({\mathbf{x}}_t,t)$。在 diffusers 库中已经实现了一个 2D UNet 网络，我们直接使用即可。下面我们也主要使用 diffusers 实现 DDPM 模型。

训练参数

首先配置训练的参数：

from dataclasses import dataclass

@dataclass
class TrainingConfig:
    image_size = 64
    train_batch_size = 16
    eval_batch_size = 16
    num_epochs = 50
    gradient_accumulation_steps = 1
    learning_rate = 1e-4
    lr_warmup_steps = 500
    mixed_precision = "fp16"
    output_dir = "ddpm-animefaces-64"
    overwrite_output_dir = True

config = TrainingConfig()

训练数据

我们使用 huggan/anime-faces 数据集，这个数据集由 21551 张分辨率为 64x64 的动漫人物头像组成。我们加载这个数据集：

from datasets import load_dataset

dataset = load_dataset("huggan/anime-faces", split="train")
dataset = dataset.select(range(21551))

由于这个数据集的作者组织数据的方式不太规范，所以最后加载进来实际上数据集的长度是 86204，也就是 21551 张图片每张重复了 4 次，我们只需要保留前 21551 个样本即可。

然后为数据集设置预处理函数：

from torchvision import transforms

def get_transform():
    preprocess = transforms.Compose([
        transforms.Resize(config.image_size),
        transforms.ToTensor(),
        transforms.Normalize([0.5], [0.5]),
    ])
    def transform(samples):
        images = [preprocess(img.convert("RGB")) for img in samples["image"]]
        return dict(images=images)
    return transform

dataset.set_transform(get_transform())

最后创建 dataloader：

from torch.utils.data import DataLoader

dataloader = DataLoader(dataset, batch_size=config.train_batch_size, shuffle=True)

降噪网络

我们可以直接用 diffusers 创建降噪网络：

from diffusers import UNet2DModel

model = UNet2DModel(
    sample_size=config.image_size,
    in_channels=3,
    out_channels=3,
    layers_per_block=2,
    block_out_channels=(128, 128, 256, 256, 512, 512),
    down_block_types=(
        "DownBlock2D",
        "DownBlock2D",
        "DownBlock2D",
        "DownBlock2D",
        "AttnDownBlock2D",
        "DownBlock2D",
    ),
    up_block_types=(
        "UpBlock2D",
        "AttnUpBlock2D",
        "UpBlock2D",
        "UpBlock2D",
        "UpBlock2D",
        "UpBlock2D",
    ),
)

核心代码

前边的三个部分分别配置了一些训练参数，以及训练数据和模型，这些都是比较工程化的部分，而我们在上面推导的 DDPM 核心算法还没有实现。在这一小节我们主要来实现核心的算法。

首先我们需要先定义 $\beta$、$\alpha$，以及 $\bar{\alpha }$ 等最基本的常量，这里我们保持 DDPM 原论文的配置，也就是 $\beta$ 初始为 $1\times 10^{-4}$，最终为 $0.02$，且共有 $1000$ 个时间步：

import torch

class DDPM:
    def __init__(
        self,
        num_train_timesteps:int = 1000,
        beta_start: float = 0.0001,
        beta_end: float = 0.02,
    ):
        self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
        self.alphas = 1.0 - self.betas
        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
        self.timesteps = torch.arange(num_train_timesteps - 1, -1, -1)

然后是比较简单的前向过程，只需要实现加噪即可，按照 ${\mathbf{x}}_t=\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ$ 这个公式实现即可。注意需要将系数的维度数量都与输入样本对齐：

class DDPM:
    ...
    
    def add_noise(
        self,
        original_samples: torch.Tensor,
        noise: torch.Tensor,
        timesteps: torch.Tensor,
    ):
        alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device ,dtype=original_samples.dtype)
        noise = noise.to(original_samples.device)
        timesteps = timesteps.to(original_samples.device)

        # \sqrt{\bar\alpha_t}
        sqrt_alpha_prod = alphas_cumprod[timesteps].flatten() ** 0.5
        while len(sqrt_alpha_prod.shape) < len(original_samples.shape):
            sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
        
        # \sqrt{1 - \bar\alpha_t}
        sqrt_one_minus_alpha_prod = (1.0 - alphas_cumprod[timesteps]).flatten() ** 0.5
        while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape):
            sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
        
        return sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise

反向过程相对来说比较复杂，不过因为我们已经完成了公式的推导，只需要按照公式实现即可。我们也再把公式贴到这里，对着公式实现具体的代码：
$$\begin{array}{rl}\sigma & ={\left(\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\bar{\alpha }}_{t-1}}\right)}^{-1/2}\\ \mu & =\frac{1}{\sqrt{{\alpha }_t}}\left({\mathbf{x}}_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\widetilde{ϵ}}_t\right)\end{array}$$

class DDPM:
    ...

    @torch.no_grad()
    def sample(
        self,
        unet: UNet2DModel,
        batch_size: int,
        in_channels: int,
        sample_size: int,
    ):
        betas = self.betas.to(unet.device)
        alphas = self.alphas.to(unet.device)
        alphas_cumprod = self.alphas_cumprod.to(unet.device)
        timesteps = self.timesteps.to(unet.device)
        images = torch.randn((batch_size, in_channels, sample_size, sample_size), device=unet.device)
        for timestep in tqdm(timesteps, desc='Sampling'):
            pred_noise: torch.Tensor = unet(images, timestep).sample

            # mean of q(x_{t-1}|x_t)
            alpha_t = alphas[timestep]
            alpha_cumprod_t = alphas_cumprod[timestep]
            sqrt_alpha_t = alpha_t ** 0.5
            one_minus_alpha_t = 1.0 - alpha_t
            sqrt_one_minus_alpha_cumprod_t = (1 - alpha_cumprod_t) ** 0.5
            mean = (images - one_minus_alpha_t / sqrt_one_minus_alpha_cumprod_t * pred_noise) / sqrt_alpha_t
            
            # variance of q(x_{t-1}|x_t)
            if timestep > 0:
                beta_t = betas[timestep]
                one_minus_alpha_cumprod_t_minus_one = 1.0 - alphas_cumprod[timestep - 1]
                one_divided_by_sigma_square = alpha_t / beta_t + 1.0 / one_minus_alpha_cumprod_t_minus_one
                variance = (1.0 / one_divided_by_sigma_square) ** 0.5
            else:
                variance = torch.zeros_like(timestep)
            
            epsilon = torch.randn_like(images)
            images = mean + variance * epsilon
        images = (images / 2.0 + 0.5).clamp(0, 1).cpu().permute(0, 2, 3, 1).numpy()
        return images

训练与推理

最后是训练和推理的代码，这部分也比较工程，直接套用现成代码即可：

from accelerate import Accelerator
from diffusers.optimization import get_cosine_schedule_with_warmup
from diffusers.utils import make_image_grid, numpy_to_pil
import torch.nn.functional as F
import os

model = model.cuda()
ddpm = DDPM()
optimizer = torch.optim.AdamW(model.parameters(), lr=config.learning_rate)
lr_scheduler = get_cosine_schedule_with_warmup(
    optimizer=optimizer,
    num_warmup_steps=config.lr_warmup_steps,
    num_training_steps=(len(dataloader) * config.num_epochs),
)
accelerator = Accelerator(
    mixed_precision=config.mixed_precision,
    gradient_accumulation_steps=config.gradient_accumulation_steps,
    log_with="tensorboard",
    project_dir=os.path.join(config.output_dir, "logs"),
)
model, optimizer, dataloader, lr_scheduler = accelerator.prepare(
    model, optimizer, dataloader, lr_scheduler
)
global_step = 0
for epoch in range(config.num_epochs):
    progress_bar = tqdm(total=len(dataloader), disable=not accelerator.is_local_main_process, desc=f'Epoch {epoch}')

    for step, batch in enumerate(dataloader):
        clean_images = batch["images"]
        # Sample noise to add to the images
        noise = torch.randn(clean_images.shape, device=clean_images.device)
        bs = clean_images.shape[0]
        # Sample a random timestep for each image
        timesteps = torch.randint(
            0, ddpm.num_train_timesteps, (bs,), device=clean_images.device,
            dtype=torch.int64
        )
        # Add noise to the clean images according to the noise magnitude at each timestep
        noisy_images = ddpm.add_noise(clean_images, noise, timesteps)
        with accelerator.accumulate(model):
            # Predict the noise residual
            noise_pred = model(noisy_images, timesteps, return_dict=False)[0]
            loss = F.mse_loss(noise_pred, noise)
            accelerator.backward(loss)
            accelerator.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()
        progress_bar.update(1)
        logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0], "step": global_step}
        progress_bar.set_postfix(**logs)
        accelerator.log(logs, step=global_step)
        global_step += 1

    if accelerator.is_main_process:
        # evaluate
        images = ddpm.sample(model, config.eval_batch_size, 3, config.image_size)
        image_grid = make_image_grid(numpy_to_pil(images), rows=4, cols=4)
        samples_dir = os.path.join(config.output_dir, 'samples')
        os.makedirs(samples_dir, exist_ok=True)
        image_grid.save(os.path.join(samples_dir, f'{global_step}.png'))
        # save models
        model.save_pretrained(config.output_dir)

结果展示

训练在一张 NVIDIA GeForce RTX 4090 GPU 上大概需要运行 3 个多小时，最后的结果大概长这个样子：

可以看到虽然里边难免有一些比较奇形怪状的结果，不过总体上来说已经初具雏形了。

总结

本文总结了 DDPM 的理论和实现方式，在代码部分我们是完全根据推导出的公式实现的采样过程。实际上在很多代码库中，采样过程并没有严格按照论文中的公式实现，而是先从 ${\mathbf{x}}_t$、$t$ 和预测的噪声反向计算出 ${\mathbf{x}}_0$，再基于 $\mu =\frac{\sqrt{{\alpha }_t}(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}{\mathbf{x}}_t+\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t}{1-{\bar{\alpha }}_t}{\mathbf{x}}_0$ 计算均值，这样的好处在于可以对 ${\mathbf{x}}_0$ 进一步规范化，控制输出的范围。

可以看出 DDPM 虽然理论比较复杂，但实现起来还是比较简单直接的。因为作者本人对 diffusion models 的理解也不算非常深入，所以如果文章有问题的话欢迎各位读者来讨论，后续（如果没有鸽掉的话）还会更新一些其他的 diffusion models 的文章，欢迎追更）

本文完整的代码以及训练好的模型见如下链接：

• 完整代码：https://github.com/LittleNyima/code-snippets/tree/master/ddpm-tutorial
• 模型权重：https://huggingface.co/LittleNyima/ddpm-anime-faces-64

参考资料：

1. 简单基础入门理解Denoising Diffusion Probabilistic Model，DDPM扩散模型
2. 扩散模型之DDPM
3. Denoising Diffusion-based Generative Modeling: Foundations and Applications
4. Train a diffusion model

本文原文以 CC BY-NC-SA 4.0 许可协议发布于 笔记｜扩散模型（一）：DDPM 理论与实现，转载请注明出处。