[Toturial] stable diffustion 扩散模型介绍扩散模型在 2015 年提出，进过 DDPM ，潜

扩散模型到 Stable Diffusion

扩散模型

扩散模型的概念在 2015 年由斯坦福大学的一博士后 Sohl-Dickstein 通过论文《Deep Unsupervised Learning using Nonequilibrium Thermodynamics (mlr.press)》提出。扩散模型的灵感来自非平衡热力学，通过定义了一个扩散步骤的马尔可夫链，以缓慢地将符合搞死分布的随机噪声添加到数据中，然后反转扩散过程以从噪声中构建所需要的数据样本

随后， 2019 年斯坦福在读博士宋彪和其导师通过论文《Generative Modeling by Estimating Gradients of the Data Distribution》提出了一种新的方法来构建生成模型，通过估计分布的梯度（分布的梯度可以看出是高维曲面的斜率），代替了之前估计数据的概率分类（数据的概率分布类似高维曲面）

DDPM

之后，2020年6月加州伯克利的 Jonathan Ho 等人基于宋彪的工作，改进了 Sohl-Dickstein 的扩散模型，通过论文《Denoising Diffusion Probabilistic Models》正式提出对于普通扩散模型的改进版： DDPM ，论文名的既是全称。 DDPM 的两大特色：

从预测转换图像改进为预测噪声，作者认为从 [Math Processing Error] 预测 [Math Processing Error] ,这种图像到图像的转化不太好优化。所以换了思路，转变成预测图像转化需要的差值，一旦预测出来，图像减去差值既得到目标图像，这个差值称做噪声。DDPM 采用了一个 U-Net 结构的 AutoEncoder 来对图像某一个时刻的噪声进 z 进行预测，模型训练目标即希望预测的噪声和真实噪声一致。
DDPM 只预测正态分布的均值。虽然正态分布有均值和方差决定，但作者发现只需求均值，逆向过程中的高斯分布的方差项直接使用一个常数，模型的效果就已经很好。

到 2021年2月，由 OpenAI 的 Alex Nichol 和 Prafulla Dhariwal 通过论文 [2102.09672] Improved Denoising Diffusion Probabilistic Models (arxiv.org) 提出了 improved DDPM。

DDPM 的逆向过程中，高斯分布的方差直接使用了一个常数， improved DDPM 的作者认为如果对方差也进行学习，效果应该会更好。
DDPM 添加噪声是采用线性的 variance schedule 改为余弦 schedule, 效果更好。
简单尝试了 scale 大模型后，生成效果更好。

2021年5月 OpenAI 的 improved DDPM 作者借鉴 Diffusion Model Beat GANs 中的 classifier guidance ，发布了《Diffusion Models Beat GANs on Image Synthesis》对 imporoved DDPM 进行改进，使用 classifier guidance 的方法，引导模型进行采样和生成，不仅生成图片更逼真，且加速了反向采样过程。所谓 classifier guided diffusion 是在反向过程训练 U-Net 的同时，也训练一个图片分类器，当采样之后，使用分类器获取图片的分类是否正确，指定扩散模型的采样和生成。

在最简单的 classifer guidance 之外，还有其他的引导方式：

CLIP guidance: 将简单的分类器换成 CLIP ，文本与图像就关联起来了，则可以利用文本引导采样和生成。
Image guidance: 利用图像特征和风格层面的引导，只需要一个 gram matrix

潜空间扩散模型

最后在2021年12月论文《High-Resolution Image Synthesis with Latent Diffusion Models》提出了潜在扩散模型，也是后续奠定stable diffusion的核心论文

为了使扩散模型在有限的计算资源上训练，并且保留它们的质量和灵活性，故首先训练了一个强大的预训练自编码器，这个自编码器所学习到的是一个潜在的空间，这个潜在的空间要比像素空间要小的多(可以简单粗暴的理解为就是一个被压缩或被降维的空间)，把扩散模型在这个潜在的空间去训练，大大的降低了对算力的要求，这也是Stable Diffusion比原装Diffusion速度快的原因

条件引导生成

2022年7月 Classifier-Free Diffusion Guidance 所谓classifier free guidance的方式(对应论文为《Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance》)，只是改变了模型输入的内容，除了 conditional 输入外(随机高斯噪声输入加引导信息)，还有 unconditional 的采样输入，两种输入都会被送到同一个 diffusion model，从而让其能够具有无条件和有条件生成的能力

扩散模型本来训练就很贵了，classifier free guidance这种方式在训练时需要生成两个输出，所以训练更贵了。但是这个方法确实效果好，所以在GLIDE 、DALL·E2和Imagen里都用了，而且都提到这是一个很重要的技巧，用了这么多技巧之后，GLIDE终于是一个很好的文生图模型了，只用了35亿参数，生成效果和分数比120亿参数的DALL·E还要好

Stable Diffusion 文生图模型介绍

于 2022 年由初创公司StabilityAI、CompVis与Runway合作开发开源版本的 Stable Diffusion 模型，支持文生图能力。

模型版本

Stable Diffusion Checkpoint 基础模型是使用 Stable Diffusion 生成图像必须具备的一个主模型。这类模型是由官方提供的，也可以用来直接绘画，但是没有固定风格，一般不拿来直接绘画训练到时候一般是将官方模型作为底模，加上自己的独特风格素材，来训练自己想要的模型官方提供的基础模型版本如下：

SD 1.1-1.4 出图分辨率为 512x512，该版本由 CompVis 提供：github.com/CompVis/sta… compvis 发布了 sd-v1-1.ckpt到sd-v1-4.ckpt 的版本
SD 1.5 出图分辨率为 512x512，runwayml 在 1.4 的基础上，推出了sd-v1-5.ckpt 的版本 github.com/runwayml/st…
SD 2.0-2.1 出图分辨率为 768x768，该版本则是stabilityai维护：github.com/Stability-A… Stability AI · GitHub
SD XL 出图分辨率为 1024x1024

以上基础模型可以到 huggingface 里面搜索组织名称StabilityAI、CompVis与Runway进行下载

1.4 CompVis/stable-diffusion-v1-4 · HF Mirror (hf-mirror.com)

1.5 runwayml/stable-diffusion-v1-5 · HF Mirror (hf-mirror.com)

2 stabilityai/stable-diffusion-2 · HF Mirror (hf-mirror.com)

2.1 stabilityai/stable-diffusion-2-1 · HF Mirror (hf-mirror.com)

vae stabilityai/sd-vae-ft-mse at main (hf-mirror.com)

模型结构

Stable Diffusion 基础模型拥有 U-Net, VAE, TextEncoder, 三部分。

U-Net 扩散模型主体，用来实现文本引导下的 Latent 生成
VAE(Variational Auto-Encoder)：Encoder 将图像编码到 latent 空间， Decoder 则将 latent 解码为图像。
TextEncoder: 提取输入 text 的 embeddings 通过 cross attentsion 的方式进入扩散模型 U-Net 中作为 Condition.

U-Net

Stable Diffusion 中的 Diffusion 是指扩散模型，因为它的数学原理看起来像物理学中的扩散而得名，扩散模型分为前向扩散和反向扩散两个过程。

前向扩散（加噪）是选择一张图像 X0 ，然后生成随机标准高斯噪声图，将标准高斯噪声图添加到到图像 X0 生成图像 X1。重复循环 t 个时间步，逐步添加经过采样的标准高斯分布噪声图，图像 X0 变化到 Xt，原始图像则被噪声完全覆盖。

前向扩散过程中的采样的标准高斯噪声图如何生成？在数学上，高斯噪声是服从均值为μ和标准差为σ的正态分布的随机值，也被称为白噪声。其概率密度函数(PDF)定义：

$pdf(x) = (1/ (σ*sqrt(2*π))) *e^(- (x—μ)²/ (2*σ²))$

其中 x 是随机变量，μ 是均值，σ 是标准差。用 𝑁(𝑥;0,1) 表示随机变量 𝑥 服从均值为0方差为1的标准正态分布，标准高斯噪声图就是从标准正态分布中采样噪声 𝜖～𝑁(0,1) 。

然后通过下面公式将标准高斯噪声的添加到图像 X0 中：

反向扩散（去噪）是从标准正态分布 𝑁(𝑥;0,1) 中随机采样一张标准高斯噪声图 Xt，然后通过预估当前 t 时刻加入的正向扩散中添加的标准高斯噪声图，从其从图像 Xt 中减去。重复循环 t 个时间步，逐步减去预估的高斯分布噪声，将随机采样的标准高斯噪声还原成原图像。

前向扩散中添加标准高斯噪声的分布均值和方差都基于 Xt 和 𝛽t 确定。反向扩散中减去的标准高斯噪声图则是通过噪声估计器 U-Net 模型预估。

什么叫U-Net模型(对应论文为：U-Net: Convolutional Networks for Biomedical Image Segmentation) U-Net 架构，骨干网络是由一系列卷积和自注意力层构成。

收缩路径遵循卷积网络的典型架构。由两个3X3卷积的重复结构组成，每层后面激活函数 ReLU 和2X2的最大池化操作，步长为 2，用于下采样。

扩展路径中的每个步骤都包括特征图的上采样，然后是一个2X2 的卷积，将特征通道数量减半，与收缩路径中相应的裁剪的特征图进行连接，以及两个 3X3 简介，每层后面激活函数 ReLU。

在最后一层，使用 1X1 卷积将每个 64 分量的特征向量映射到所需的类别数。

根据上图(注意该图下方是前向过程，上方是反向过程)前向过程加入噪声的分布是已知的正如 𝑞(𝑥𝑡|𝑥𝑡−1) ，但是 𝑝𝜃（𝑥𝑡−1|𝑥𝑡）的反向过程的高斯噪声是未知的。 U-Net 模型就是噪声估计模型 𝜖(Xt, t) ,使其预估的噪声与前向过程中的标准高斯噪声相近，具体训练过程：

选择一个训练图像，执行前向扩散的过程。
通过 U-Net 模型预测添加的噪声，使其与前向扩散中的实际添加的标准高斯噪声做比较，从而建立损失函数做反向传播，更新噪声估计器的参数。在训练过程结束后我们就可以得到一个模型可以预测 𝑡 时刻加入的标准高斯分布噪声 𝜖 。

Stable Diffusion训练数据集的具体情况，Stable Diffusion的训练是多阶段的（先在256×256尺寸上预训练，然后在512×512尺寸上精调），不同的阶段产生了不同的版本：

SD v1.1：在Laion2B-en数据集上以256×256大小训练237,000步，如上所述，laion2B-en数据集中256以上的样本量共1324M；然后在Laion5B的高分辨率数据集以512×512尺寸训练194,000步，这里的高分辨率数据集是图像尺寸在1024×1024以上，共170M样本。

SD v1.2：以SD v1.1为初始权重，在improved_aesthetics_5plus数据集上以512×512尺寸训练515,000步数，这个improved_aesthetics_5plus数据集上laion2B-en数据集中美学评分在5分以上的子集（共约600M样本），注意这里过滤了含有水印的图片（pwatermark>0.5)以及图片尺寸在512×512以下的样本。

SD v1.3：以SD v1.2为初始权重，在improved_aesthetics_5plus数据集上继续以512×512尺寸训练195,000步数，不过这里采用了CFG（以10%的概率随机drop掉text）。

SD v1.4：以SD v1.2为初始权重，在improved_aesthetics_5plus数据集上采用CFG以512×512尺寸训练225,000步数。 *SD v1.5：以SD v1.2为初始权重，在improved_aesthetics_5plus数据集上采用CFG以512×512尺寸训练595,000步数。

VAE

Stable Diffusion 原来名字叫做 "Latent Diffusion Model", 是基于 Latent 的扩散模型，像 Imagen, DALL-E 等扩散模型是基于像素空间工作。其中的 Latent 则是潜空间。它不在高维的图像空间中操作,而是先将图像压缩到潜变量空间。潜变量空间的维度是图像空间的48倍小,所以它可以大大减少计算量。这就是为什么它的速度会快很多。

Stable Diffusion 是基于 latent 潜空间的扩散模型，潜空间是使用 VAE(Variational Auto-Encoder) 变分自编码编码器将图片映射到一个低纬的的隐式表征。在潜空间进行扩散可以使用更少的内存，减少 U-Net 层数，大大的减少计算量，加速图片的生成。

VAE(Variational Auto-Encoder)是一种神经网络模型的一部分。 VAE 提供 Stable Diffusion 对图像进行潜空间与像素空间之间的转化时需要的编码与解码器。

一般无需安装 VAE 文件就可以运行 Stable Diffusion,因为无论是 v1、v2 还是自定义模型,都已经内置了默认的 VAE。
当需要下载和使用 VAE 时,通常是需要使用 VAE 的改进版本。这发生在模型训练者用额外的数据进一步微调了模型的 VAE 部分。为了避免发布一个很大的完整新模型,他们只发布了经过更新的小部分。

Stability AI发布了两种精调后的VAE解码器变体,ft-EMA和ft-MSE，它们强调的部分不同！ EMA和MSE：指数移动平均（Exponential Moving Average）和均方误差（Mean Square Error）是测量自动编码器好坏的指标。

TextEncoder

Stable Diffusion 中文本 prompt 的 embedding 使用了 CLIP 模型，CLIP是一个由OpenAI开发的深度学习模型,用于为任何图像生成文本描述。Stable Diffusion v1使用了CLIP的标记器。

Stable Diffusion 是文生图模型，所以在扩散模型生图的过程中可以使用文本 prompt 来控制图像的生成。文本 prompt 进行 embedding 获得相应的嵌入向量，输入到扩散模型中，进行 cross-attention 跨注意力计算，从而影响生成的图像。

《扩散模型-从原理到实战》

环境准备

安装显卡驱动与 CUDA
安装机器学习开源库 torch
安装 HuggingFace 的扩散模型开源库 diffusers
安装绘图开源库 matplotlib

使用以下代码进行环境和数据集测试

import torch
import torchvision
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader
from diffusers import DDPMScheduler, UNet2DModel
from matplotlib import pyplot as plt

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f'Using device: {device}')

dataset = torchvision.datasets.MNIST(root="mnist/", train=True, download=True, transform=torchvision.transforms.ToTensor())
train_dataloader = DataLoader(dataset, batch_size=8, shuffle=True)
x, y = next(iter(train_dataloader))
print('Input shape:', x.shape)
print('Labels: ', y)
plt.imshow(torchvision.utils.make_grid(x)[0], cmap='Greys')
plt.show()

退化过程

如果 amount = 0, 则返回输入，不做任何更改；如果 amount = 1, 将得到一个纯粹的噪声。通过控制 amount 的值，可以将输入内容与噪声混合，并把混合后结果保持在相同的范围(0-1)

def corrupt(x, amount):
    """ 根据 amount 为输入 x 加入噪声，这就是退化的过程 """
    noise = torch.rand_like(x)
    amount = amount.view(-1, 1, 1, 1) # 整理形状以保证广播机制不出错
    return (1 - amount) * x + amount * noise

# 绘制输入数据
fig, axs = plt.subplots(2, 1, figsize=(12, 5))
axs[0].set_title('Input data')
axs[0].imshow(torchvision.utils.make_grid(x)[0], cmap='Greys')

# 加入噪声
amount = torch.linspace(0, 1, x.shape[0]) # 从 0 到 1 退化更强烈了
noised_x = corrupt(x, amount)

# 绘制加噪版本的图像
axs[1].set_title('Corrupted data (--amount increased -->)')
axs[1].imshow(torchvision.utils.make_grid(noised_x)[0], cmap='Greys')

plt.show()

torch.linspace(start, end, steps) : 全称 linear space 线性等分向量。构造一个一维向量，值是从 start 到 end 之间，基于 steps 等分。同时向量 size 为 steps.

torch.rand_like(input): 返回一个与 input 相同 size 的向量，填充值从 0-1 ，且符合均匀分布的随机值。

torchvision.utils.make_grid(tensor: Union[Tensor, List[Tensor]]): 将 tensor 中包含的图片组合成一个包含所有图片的 grid 图片

模型训练

U-Net 网络

U-Net 网络最初用于完成医学图像中的分割任务，在扩散模型中用于接受 28X28 像素的噪声图像，并输出相同大小图片的预测结果。U-Net 网络由一条“压缩路径”和一条“扩展路径” 组成。“压缩路径”会使通过该路径的数据维度压被压缩，而“扩展路径”则会将数据扩展回原始维度。网络中的残差连接允许信息和梯度在不同层级之间流动。

在这里，仅构建一个非常简单的示例，接收一个单通道的图像，使其通过压缩路径的 3 个卷积层和扩展路径的 3 个卷积层。压缩路径和扩展路径之间有残差连接，使用最大池化层进行下采样

池化层：池化层（Pooling Layer）是卷积神经网络（CNN）中的一个重要组件，它主要用于降低特征的空间维度，从而减少计算量和参数数量，同时使特征检测更加鲁棒。池化层对输入的特征图进行压缩，提取主要特征，并且增加对图像位移的不变性，这包括平移不变性、旋转不变性和尺度不变性。

池化操作不涉及权重的学习，它是一个确定性的操作，通常计算的是池化窗口内的最大值或平均值，这就是所谓的最大池化（Max Pooling）和平均池化（Average Pooling）。最大池化通过取区域内的最大值来实现特征的下采样，而平均池化则是计算区域内的平均值。

池化层的作用主要包括：

特征不变性的提升，使模型更加关注特征是否存在而不是具体位置。

特征降维，减小下一层的输入大小，降低计算量和参数个数。

防止过拟合，方便模型优化。

实现非线性，类似于ReLU激活函数。

扩大感受野，使网络能够捕捉更广泛的特征。

常见的池化层包括最大池化层、平均池化层、重叠池化层以及其他形式的池化层。重叠池化层允许相邻池化窗口之间有重叠区域，这可以通过设置池化窗口大小和步幅（Stride）来实现，通常窗口大小大于步幅。

在实际应用中，池化层通常跟在卷积层和激活函数之后，帮助网络逐步聚合信息，形成更加全局的特征表示，同时保持卷积层的优势。池化层没有模型参数，因此不需要进行参数初始化或更新。

池化层的实现通常很简单，例如在PyTorch中，可以通过nn.MaxPool2d或nn.AvgPool2d来实现最大池化或平均池化，并且可以通过设置padding和stride参数来调整输出特征图的大小。

import torch
import torchvision
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader
from diffusers import DDPMScheduler, UNet2DModel
from matplotlib import pyplot as plt

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f'Using device: {device}')

class BasicUNet(nn.Module):
    """一个十分简单的 U-Net 网络部署"""
    def __init__(self, in_channels=1, out_channels=1):
        super().__init__()
        self.down_layers = torch.nn.ModuleList([
            nn.Conv2d(in_channels, 32, kernel_size=5, padding=2),
            nn.Conv2d(32, 64, kernel_size=5, padding=2),
            nn.Conv2d(64, 64, kernel_size=5, padding=2)
        ])
        self.up_layers = torch.nn.ModuleList([
            nn.Conv2d(64, 64, kernel_size=5, padding=2),
            nn.Conv2d(64, 32, kernel_size=5, padding=2),
            nn.Conv2d(32, out_channels, kernel_size=5, padding=2)
        ])
        self.act = nn.ReLU() # 激活函数
        self.downscale = nn.MaxPool2d(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: # 选择除了第 3 层（最后一层）以外的层
                h.append(x) # 排列供残差连接使用的数据
                x = self.downscale(x) # 进行下采样以适配下一层的输入

        for i, l in enumerate(self.up_layers):
            if i > 0: # 选择除了第 1 个上采样层以外的层
                x = self.upscale(x) # Upscale 上采样
                x += h.pop() # 得导之前排列好的供残差连接使用的数据
            x = self.act(l(x)) # 通过运算层与激活函数

        return x
    
if __name__ == '__main__':
    net = BasicUNet()
    x = torch.rand(8, 1, 28, 28)
    print("shape: ", net(x).shape)

    print("param size: ", sum([p.numel() for p in net.parameters()]))

激活函数是神经网络中不可或缺的组成部分，它为模型引入非线性，使得网络能够学习和执行更加复杂的任务。以下是一些常见的激活函数：

Sigmoid：这是一个将输入压缩到0和1之间的函数，通常用于二分类问题。

Tanh (Hyperbolic Tangent) ：这是双曲正切函数，将输入压缩到-1和1之间。

ReLU (Rectified Linear Unit) ：当前最流行的激活函数之一，它将所有负值置为0，保留正值。公式为：ReLU(x)=max⁡(0,x)。它计算简单，训练速度快，且在很多任务中表现良好。

Leaky ReLU：这是ReLU的变体，允许小的梯度值当输入为负数时。公式为：Leaky ReLU(x)=max⁡(αx,x)，其中αα是一个很小的正数。

Parametric ReLU (PReLU) ：Leaky ReLU的泛化形式，其斜率系数αα是可学习的参数。

ELU (Exponential Linear Unit) ：它在负值域内以指数速率衰减，同时保持正值不变。公式为：ELU(x)=max⁡(α(ex−1),x)。

SELU (Scaled Exponential Linear Unit) ：自归一化激活函数，它不仅自归一化，而且输出的均值和标准差在训练过程中保持恒定。公式为：SELU(x)=λ(max⁡(0,x)−αe−max⁡(0,−x))。

Softmax：通常用于多分类问题的最后一层，将输入的任意实数值转换成概率分布。

Swish：一种自门控的激活函数，公式为：f(x)=x⋅σ(βx)，其中β是可学习的参数或一个常数。

GELU (Gaussian Error Linear Unit) ：基于高斯分布累积分布函数的激活函数，近期在自然语言处理领域流行。公式为：GELU(x)=xΦ(x)，其中Φ是标准正态分布的CDF。

激活函数的选择取决于具体的应用和网络架构。例如，ReLU及其变体由于其计算效率和在深层网络中的效果，成为许多现代网络的默认选择。然而，对于一些需要输出概率或需要梯度在负值域内传播的任务，可能需要使用如Softmax或GELU这样的激活函数。

开始训练

针对原始输入 x，首先给定一个“带噪”的输入 noisy_x ，扩散模型输出预测结果。通过均方差对预测结果与原始值进行比较。流程如下：

获取一批训练数据
添加随机噪声
将数据输入扩散模型
对模型输出预测结果与原始图像进行比较，计算损失，反向传播更新模型参数

训练过程中的损失曲线

模型预测结果可视化展示


# 数据加载器（可以调整 batch_size)
batch_size = 128
dataset = torchvision.datasets.MNIST(root="mnist/", train=True, download=True, transform=torchvision.transforms.ToTensor())
train_dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

# 设置将在整个数据集上运行多个个周期
n_epochs = 3

# 创建网络
net = BasicUNet()
net.to(device)

# 指定损失函数
loss_fn = nn.MSELoss()

# 指定优化器
opt = torch.optim.Adam(net.parameters(), lr = 1e-3)

# 记录训练过程中的损失，供后期查看
losses = []


# 训练
for epoch in range(n_epochs):
    for x, y in train_dataloader:
        # 得到数据并准备退化
        x = x.to(device) # 将数据加载到 GPU 显存中
        noise_amount = torch.rand(x.shape[0]).to(device) # 随机选取噪声量
        noisy_x = corrupt(x, noise_amount) # 创建“带噪”的输入 noisy_x

        # 得到模型预测结果
        pred = net(noisy_x)

        # 计算损失函数
        loss = loss_fn(pred, x) # 输出与真实“干净”的 x 有多接近

        # 反向传播并更新参数
        opt.zero_grad()
        loss.backward()
        opt.step()

        # 存储损失，供后期查看
        losses.append(loss.item())

    # 输出在每个周期训练得到的损失的均值
    avg_loss = sum(losses[-len(train_dataloader):])/len(train_dataloader)
    print(f'Finished epoch {epoch}. Average loss for this epoch: {avg_loss:05f}')


# 查看损失曲线
plt.plot(losses)
plt.ylim(0, 0.1)
plt.show()

# 可视化模型在“带噪”输入上的表现
# 构造测试集
x, y = next(iter(train_dataloader))
x = x[:8] # 只选择前 8 条数据

# 在 (0, 1) 区间选择退化量
amount = torch.linspace(0, 1, x.shape[0]) # 从 0 到 1 退化越强烈
noised_x = corrupt(x, amount)

# 得到模型预测结果
with torch.no_grad():
    preds = net(noised_x.to(device)).detach().cpu()

# 绘制结果
fig, axs = plt.subplots(3, 1, figsize=(12, 7))
axs[0].set_title('Input data')
axs[0].imshow(torchvision.utils.make_grid(x)[0].clip(0, 1), cmap='Greys')

axs[1].set_title('Corrupted data')
axs[1].imshow(torchvision.utils.make_grid(noised_x)[0].clip(0, 1), cmap='Greys')

axs[2].set_title('Network Predictions data')
axs[2].imshow(torchvision.utils.make_grid(preds)[0].clip(0, 1), cmap='Greys')

plt.show()

nn.MSELoss(): 创建一个损失函数，用来衡量预测结果与目标结果之间的差值。

损失函数（Loss Function），也称为代价函数或目标函数，是衡量模型预测值与实际值差异的函数。在训练神经网络时，目标是最小化损失函数，以此来调整网络的参数。以下是一些常见的损失函数：

均方误差（Mean Squared Error, MSE） ：这是回归问题中最常用的损失函数。它计算预测值与实际值之差的平方的平均值。

平均绝对误差（Mean Absolute Error, MAE） ：与MSE类似，但计算的是绝对值的平均，对异常值的敏感度较低。

交叉熵损失（Cross-Entropy Loss） ：用于分类问题，特别是二分类和多分类问题。对于二分类问题。

对数损失（Logarithmic Loss） ：与交叉熵损失相同，常用于二分类问题。

Hinge损失（Hinge Loss） ：用于支持向量机（SVMs），计算间隔的非负性。

Categorical Cross-Entropy Loss：多分类问题的交叉熵损失，用于神经网络的最后层，当进行多类别分类时。

Dice损失（Dice Loss） ：用于计算两个样本之间的相似度，特别适用于医学图像分割。

Focal Loss：由何凯明等人提出，用于解决类别不平衡问题，特别是用于目标检测和分类。

IoU损失（Intersection over Union Loss） ：用于计算预测的边界框和真实边界框之间的重叠程度，常用于目标检测。

三元组损失（Triplet Loss） ：用于训练具有三元组的模型，包含一个锚点、一个正样本和一个负样本。目标是使锚点与正样本的距离小于锚点与负样本的距离。

不同的损失函数适用于不同类型的问题和模型。选择合适的损失函数对于训练有效的神经网络至关重要。

采样过程

上面的模型预测结果可视化展示，可以看到噪声量越多，预测结果越差。那么从完全随机的噪声开始，一次预测是很难获取好的结果。扩散模型的优化方式则是通过多步预测，如果新的一次的预测结果比上一次的更好一点，那么将每次预测的结果与本次的输入进行融合，作为下一次预测的输入，那么我们就持续朝着更好的方向移动。

下图则是分 5 步进行预测的过程，左侧是 5 次的输入值，右侧是 5 次的预测结果，可以看到结果越往后越好。

这中分步的过程称做采样，采样的次数越多，模型推理获得的图像质量越高。另外模型训练的 epoch 越大，调整学习率，优化器等，以及使用更大的训练数据集，获得更好的模型也可以优化预测图像的质量。

# 把采样过程拆解为 5 步，每次只前进一步
n_steps = 5
x = torch.rand(8, 1, 28, 28).to(device) # 符合正态分布的随机值
step_history = [x.detach().cpu()]
pred_output_history = []

for i in range(n_steps):
    with torch.no_grad(): # 在推理的时候步需要考虑张量的导数
        pred = net(x)
    pred_output_history.append(pred.detach().cpu()) # 将模型的预测输出保存下来，以便后续绘制

    mix_factor = 1/(n_steps - i) # 设置朝着预测方向移动多少
    x = x * (1 - mix_factor) + pred * mix_factor
    step_history.append(x.detach().cpu())


# 绘制结果
fig, axs = plt.subplots(n_steps, 2, figsize=(9, 4), sharex=True)
axs[0,0].set_title('x (model Input data)')
axs[0,1].set_title('model prediction data')

for i in range(n_steps):
    axs[i, 0].imshow(torchvision.utils.make_grid(step_history[i])[0].clip(0, 1), cmap='Greys')
    axs[i, 1].imshow(torchvision.utils.make_grid(pred_output_history[i])[0].clip(0, 1), cmap='Greys')

plt.show()

DDPM

Diffusers 库中 DDPM 版本实现过程的区别：

UNet2DModel 模型结构相比 BasicUNet 模型结构更优。
退化过程的处理方式更复杂
训练目标不同，旨在预测噪声，不是“去噪”后的图像
UNet2DModel 模型通过调节时间步来调节噪声量，时间 t 作为一个额外参数被传入前向过程。
有更多种类的采样策略可供选择。

Finished epoch 0. Average loss for this epoch: 0.020083
Finished epoch 1. Average loss for this epoch: 0.012908
Finished epoch 2. Average loss for this epoch: 0.011645


from diffusers import UNet2DModel


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f'Using device: {device}')

def corrupt(x, amount):
    """ 根据 amount 为输入 x 加入噪声，这就是退化的过程 """
    noise = torch.rand_like(x)
    amount = amount.view(-1, 1, 1, 1) # 整理形状以保证广播机制不出错
    return (1 - amount) * x + amount * noise

# 数据加载器（可以调整 batch_size)
batch_size = 128
dataset = torchvision.datasets.MNIST(root="mnist/", train=True, download=True, transform=torchvision.transforms.ToTensor())
train_dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

# 设置将在整个数据集上运行多个个周期
n_epochs = 3

# 创建网络
net = UNet2DModel(
    sample_size=28,
    in_channels=1,
    out_channels=1,
    layers_per_block=2,
    block_out_channels=(32, 64, 64),
    down_block_types=(
        "DownBlock2D",
        "AttnDownBlock2D",
        "AttnDownBlock2D",
    ),
    up_block_types=(
        "AttnUpBlock2D",
        "AttnUpBlock2D",
        "UpBlock2D",
    )
)
net.to(device)

# 指定损失函数
loss_fn = nn.MSELoss()

# 指定优化器
opt = torch.optim.Adam(net.parameters(), lr = 1e-3)

# 记录训练过程中的损失，供后期查看
losses = []


# 训练
for epoch in range(n_epochs):
    for x, y in train_dataloader:
        # 得到数据并准备退化
        x = x.to(device) # 将数据加载到 GPU 显存中
        noise_amount = torch.rand(x.shape[0]).to(device) # 随机选取噪声量
        noisy_x = corrupt(x, noise_amount) # 创建“带噪”的输入 noisy_x

        # 得到模型预测结果，这里传递 t=0, 以表明模型是在没有时间步的情况下工作，也可以输入 noise_amount * 1000, 以使时间步与噪声水平相当
        pred = net(noisy_x, 0).sample

        # 计算损失函数
        loss = loss_fn(pred, x) # 输出与真实“干净”的 x 有多接近

        # 反向传播并更新参数
        opt.zero_grad()
        loss.backward()
        opt.step()

        # 存储损失，供后期查看
        losses.append(loss.item())

    # 输出在每个周期训练得到的损失的均值
    avg_loss = sum(losses[-len(train_dataloader):])/len(train_dataloader)
    print(f'Finished epoch {epoch}. Average loss for this epoch: {avg_loss:05f}')


# Plot losses and some samples
fig, axs = plt.subplots(1, 2, figsize=(12, 5))

# 查看损失曲线
axs[0].plot(losses)
axs[0].set_ylim(0, 0.1)
axs[0].set_title('Loss over time')


# 把采样过程拆解为 40 步，每次只前进一步
n_steps = 40
x = torch.rand(64, 1, 28, 28).to(device) # 符合正态分布的随机值
step_history = [x.detach().cpu()]
pred_output_history = []

for i in range(n_steps):
    noise_amount = torch.ones((x.shape[0], )).to(device) * (1 - i/n_steps)
    with torch.no_grad(): # 在推理的时候步需要考虑张量的导数
        pred = net(x, 0).sample
    pred_output_history.append(pred.detach().cpu()) # 将模型的预测输出保存下来，以便后续绘制

    mix_factor = 1/(n_steps - i) # 设置朝着预测方向移动多少
    x = x * (1 - mix_factor) + pred * mix_factor
    step_history.append(x.detach().cpu())


axs[1].imshow(torchvision.utils.make_grid(x.detach().cpu(), nrow=8)[0].clip(0, 1), cmap='Greys')
axs[1].set_title('Generated Samples')
plt.show()

退化过程

模型训练

参考

GitHub - s9roll7/animatediff-cli-prompt-travel: animatediff prompt travel

GitHub - mix1009/sdwebuiapi: Python API client for AUTOMATIC1111/stable-diffusion-webui

一文学会Stable Diffusion模型原理轻松入门AIGC - 知乎 (zhihu.com)

大白话AI | 图像生成模型DDPM | 扩散模型 | 生成模型 | 概率扩散去噪生成模型_哔哩哔哩_bilibili

AI绘画原理解析：从CLIP、BLIP到DALLE、DALLE 2、DALLE 3、Stable Diffusion_dalle3算法-CSDN博客