昇思 25 天学习打卡营第 21 天 | Diffusion扩散模型

今天是参加昇思打卡营的第21天，学习内容是Diffusion扩散模型。

以下是内容关键点概要:

1. 扩散模型简介：

   • 扩散模型是一种生成模型，通过逐步添加噪声的方式将数据从简单分布转换为复杂数据分布。
   • 逆向过程通过学习去噪，逆转前向过程，生成数据样本。

2. 模型结构：

   • 扩散模型包括正向扩散过程和逆向去噪过程。
   • 前向过程逐步添加噪声，逆向过程逐步去除噪声。

3. 正向扩散过程：

   • 将高斯噪声逐步添加到图像中，直到最终得到纯噪声。
   • 过程由时间步长 𝑡t 索引，每个时间步长 𝑡t 只与前一个时间步长相关。

4. 逆向去噪过程：

   • 学习神经网络 𝑝𝜃pθ​ 来近似条件概率分布，从纯噪声开始逐渐去噪。
   • 神经网络学习表示条件概率分布的平均值 𝜇𝜃μθ​。

5. U-Net神经网络：

   • U-Net结构被用于扩散模型，具有编码器和解码器，中间有瓶颈层。
   • 编码器和解码器之间引入了残差连接，改善梯度流。

6. 位置向量：

   • 使用正弦位置嵌入来编码时间步长 𝑡t，神经网络“知道”它在哪个特定时间步长上运行。

7. Attention模块：

   • 引入了Transformer架构中的Attention机制，提升模型性能。
   • 提供了常规的multi-head self-attention和LinearAttention两种变体。

8. 数据准备与处理：

   • 数据集通过随机水平翻转、调整大小等操作进行预处理。
   • 图像值被线性缩放到 [−1,1][−1,1] 范围内。

9. 训练过程：

   • 定义动态学习率和优化器，使用前向扩散过程生成噪声图像。
   • 训练神经网络预测噪声，优化模型以最小化损失。

10. 推理过程（采样）：

    • 从训练好的模型中生成新图像，通过逆向扩散过程从纯噪声开始，逐步去噪。

import math  # 导入数学库，用于数学运算
from functools import partial  # 导入partial函数，用于函数的部分参数调用

# 使matplotlib生成的图形能够在Jupyter Notebook中内联显示
%matplotlib inline
import matplotlib.pyplot as plt  # 导入matplotlib的pyplot模块，用于数据可视化
from tqdm.auto import tqdm  # 导入tqdm模块，用于显示进度条
import numpy as np  # 导入NumPy库，用于数学运算和处理数据结构
from multiprocessing import cpu_count  # 导入cpu_count，用于获取CPU核心数

# 导入download函数，用于下载数据集
from download import download

import mindspore as ms  # 导入MindSpore库，并设置别名为ms
import mindspore.nn as nn  # 从MindSpore库中导入nn模块，包含构建神经网络所需的类
import mindspore.ops as ops  # 从MindSpore库中导入ops模块，包含执行基础操作的函数
from mindspore import Tensor, Parameter  # 从MindSpore库中导入Tensor和Parameter类
from mindspore import dtype as mstype  # 从MindSpore库中导入dtype模块，用于数据类型转换
from mindspore.dataset.vision import Resize, Inter, CenterCrop, ToTensor, RandomHorizontalFlip, ToPIL  # 从MindSpore库中导入数据集预处理相关的类
from mindspore.common.initializer import initializer  # 从MindSpore库中导入initializer模块，用于参数初始化
from mindspore.amp import DynamicLossScaler  # 从MindSpore库中导入DynamicLossScaler，用于混合精度训练

# 设置随机种子，确保结果可复现
ms.set_seed(0)

模型简介

什么是Diffusion Model？

如果将Diffusion与其他生成模型（如Normalizing Flows、GAN或VAE）进行比较，它并没有那么复杂，它们都将噪声从一些简单分布转换为数据样本，Diffusion也是从纯噪声开始通过一个神经网络学习逐步去噪，最终得到一个实际图像。 Diffusion对于图像的处理包括以下两个过程：

• 我们选择的固定（或预定义）正向扩散过程 𝑞𝑞 ：它逐渐将高斯噪声添加到图像中，直到最终得到纯噪声

• 一个学习的反向去噪的扩散过程 𝑝𝜃𝑝𝜃 ：通过训练神经网络从纯噪声开始逐渐对图像去噪，直到最终得到一个实际的图像

由 𝑡𝑡 索引的正向和反向过程都发生在某些有限时间步长 𝑇𝑇（DDPM作者使用 𝑇=1000𝑇=1000）内。从𝑡=0𝑡=0开始，在数据分布中采样真实图像 𝐱0𝑥0（本文使用一张来自ImageNet的猫图像形象的展示了diffusion正向添加噪声的过程），正向过程在每个时间步长 𝑡𝑡 都从高斯分布中采样一些噪声，再添加到上一个时刻的图像中。假定给定一个足够大的 𝑇𝑇 和一个在每个时间步长添加噪声的良好时间表，您最终会在 𝑡=𝑇𝑡=𝑇 通过渐进的过程得到所谓的各向同性的高斯分布。

扩散模型实现原理

Diffusion 前向过程

所谓前向过程，即向图片上加噪声的过程。虽然这个步骤无法做到图片生成，但这是理解diffusion model以及构建训练样本至关重要的一步。 首先我们需要一个可控的损失函数，并运用神经网络对其进行优化。

设 𝑞(𝑥0)𝑞(𝑥0) 是真实数据分布，由于 𝑥0∼𝑞(𝑥0)𝑥0∼𝑞(𝑥0) ，所以我们可以从这个分布中采样以获得图像 𝑥0𝑥0 。接下来我们定义前向扩散过程 𝑞(𝑥𝑡|𝑥𝑡−1)𝑞(𝑥𝑡|𝑥𝑡−1) ，在前向过程中我们会根据已知的方差 0<𝛽1<𝛽2<...<𝛽𝑇<10<𝛽1<𝛽2<...<𝛽𝑇<1 在每个时间步长 t 添加高斯噪声，由于前向过程的每个时刻 t 只与时刻 t-1 有关，所以也可以看做马尔科夫过程：

                               

回想一下，正态分布（也称为高斯分布）由两个参数定义：平均值 𝜇𝜇 和方差 𝜎2≥0𝜎2≥0 。基本上，在每个时间步长 𝑡𝑡 处的产生的每个新的（轻微噪声）图像都是从条件高斯分布中绘制的，其中

                                    

我们可以通过采样 然后设置

                                      

请注意， 𝛽𝑡𝛽𝑡 在每个时间步长 𝑡𝑡 （因此是下标）不是恒定的：事实上，我们定义了一个所谓的“动态方差”的方法，使得每个时间步长的 𝛽𝑡𝛽𝑡 可以是线性的、二次的、余弦的等（有点像动态学习率方法）。

因此，如果我们适当设置时间表，从 𝐱0𝑥0 开始，我们最终得到 𝐱1,...,𝐱𝑡,...,𝐱𝑇𝑥1,...,𝑥𝑡,...,𝑥𝑇，即随着 𝑡𝑡 的增大 𝐱𝑡𝑥𝑡 会越来越接近纯噪声，而 𝐱𝑇𝑥𝑇 就是纯高斯噪声。

那么，如果我们知道条件概率分布 𝑝(𝐱𝑡−1|𝐱𝑡)𝑝(𝑥𝑡−1|𝑥𝑡) ，我们就可以反向运行这个过程：通过采样一些随机高斯噪声 𝐱𝑇𝑥𝑇，然后逐渐去噪它，最终得到真实分布 𝐱0𝑥0 中的样本。但是，我们不知道条件概率分布 𝑝(𝐱𝑡−1|𝐱𝑡)𝑝(𝑥𝑡−1|𝑥𝑡) 。这很棘手，因为需要知道所有可能图像的分布，才能计算这个条件概率。

Diffusion 逆向过程

为了解决上述问题，我们将利用神经网络来近似（学习）这个条件概率分布 𝑝𝜃(𝐱𝑡−1|𝐱𝑡)𝑝𝜃(𝑥𝑡−1|𝑥𝑡) , 其中 𝜃𝜃 是神经网络的参数。如果说前向过程(forward)是加噪的过程，那么逆向过程(reverse)就是diffusion的去噪推断过程，而通过神经网络学习并表示 𝑝𝜃(𝐱𝑡−1|𝐱𝑡)𝑝𝜃(𝑥𝑡−1|𝑥𝑡) 的过程就是Diffusion 逆向去噪的核心。

现在，我们知道了需要一个神经网络来学习逆向过程的（条件）概率分布。我们假设这个反向过程也是高斯的，任何高斯分布都由2个参数定义：

• 由 𝜇𝜃𝜇𝜃 参数化的平均值

• 由 𝜇𝜃𝜇𝜃 参数化的方差

综上，我们可以将逆向过程公式化为

                          

其中平均值和方差也取决于噪声水平 𝑡𝑡 ，神经网络需要通过学习来表示这些均值和方差。

• 注意，DDPM的作者决定保持方差固定，让神经网络只学习（表示）这个条件概率分布的平均值 𝜇𝜃𝜇𝜃 。

• 本文我们同样假设神经网络只需要学习（表示）这个条件概率分布的平均值 𝜇𝜃𝜇𝜃 。

为了导出一个目标函数来学习反向过程的平均值，作者观察到 𝑞𝑞 和 𝑝𝜃𝑝𝜃 的组合可以被视为变分自动编码器(VAE)。因此，变分下界（也称为ELBO）可用于最小化真值数据样本 𝐱0𝑥0 的似然负对数（有关ELBO的详细信息，请参阅VAE论文（Kingma等人，2013年）），该过程的ELBO是每个时间步长的损失之和 𝐿=𝐿0+𝐿1+...+𝐿𝑇𝐿=𝐿0+𝐿1+...+𝐿𝑇 ，其中，每项的损失 𝐿𝑡𝐿𝑡 （除了 𝐿0𝐿0 ）实际上是2个高斯分布之间的KL发散，可以明确地写为相对于均值的L2-loss!

如Sohl-Dickstein等人所示，构建Diffusion正向过程的直接结果是我们可以在条件是 𝐱0𝑥0 （因为高斯和也是高斯）的情况下，在任意噪声水平上采样 𝐱𝑡𝑥𝑡 ，而不需要重复应用 𝑞𝑞 去采样 𝐱𝑡𝑥𝑡 ，这非常方便。使用

                                            

我们就有

                              

这意味着我们可以采样高斯噪声并适当地缩放它，然后将其添加到 𝐱0𝑥0 中，直接获得 𝐱𝑡𝑥𝑡 。

请注意，𝛼¯𝑡𝛼¯𝑡 是已知 𝛽𝑡𝛽𝑡 方差计划的函数，因此也是已知的，可以预先计算。这允许我们在训练期间优化损失函数 𝐿𝐿 的随机项。或者换句话说，在训练期间随机采样 𝑡𝑡 并优化 𝐿𝑡𝐿𝑡 。

正如Ho等人所展示的那样，这种性质的另一个优点是可以重新参数化平均值，使神经网络学习（预测）构成损失的KL项中噪声的附加噪声。这意味着我们的神经网络变成了噪声预测器，而不是（直接）均值预测器。其中，平均值可以按如下方式计算：

                   

最终的目标函数 𝐿𝑡𝐿𝑡 如下 （随机步长 t 由 (𝜖∼𝑁(0,𝐈))(𝜖∼𝑁(0,𝐼)) 给定）：

          

在这里， 𝐱0𝑥0 是初始（真实，未损坏）图像， 𝜖𝜖 是在时间步长 𝑡𝑡 采样的纯噪声，𝜖𝜃(𝐱𝑡,𝑡)𝜖𝜃(𝑥𝑡,𝑡)是我们的神经网络。神经网络是基于真实噪声和预测高斯噪声之间的简单均方误差（MSE）进行优化的。

训练算法现在如下所示：

换句话说：

• 我们从真实未知和可能复杂的数据分布中随机抽取一个样本 𝑞(𝐱0)𝑞(𝑥0)

• 我们均匀地采样11和𝑇𝑇之间的噪声水平𝑡𝑡（即，随机时间步长）

• 我们从高斯分布中采样一些噪声，并使用上面定义的属性在 𝑡𝑡 时间步上破坏输入

• 神经网络被训练以基于损坏的图像 𝐱𝑡𝑥𝑡 来预测这种噪声，即基于已知的时间表 𝐱𝑡𝑥𝑡 上施加的噪声

实际上，所有这些都是在批数据上使用随机梯度下降来优化神经网络完成的。

U-Net神经网络预测噪声

神经网络需要在特定时间步长接收带噪声的图像，并返回预测的噪声。请注意，预测噪声是与输入图像具有相同大小/分辨率的张量。因此，从技术上讲，网络接受并输出相同形状的张量。那么我们可以用什么类型的神经网络来实现呢？

这里通常使用的是非常相似的自动编码器，您可能还记得典型的"深度学习入门"教程。自动编码器在编码器和解码器之间有一个所谓的"bottleneck"层。编码器首先将图像编码为一个称为"bottleneck"的较小的隐藏表示，然后解码器将该隐藏表示解码回实际图像。这迫使网络只保留bottleneck层中最重要的信息。

在模型结构方面，DDPM的作者选择了U-Net，出自（Ronneberger et al.，2015）（当时，它在医学图像分割方面取得了最先进的结果）。这个网络就像任何自动编码器一样，在中间由一个bottleneck组成，确保网络只学习最重要的信息。重要的是，它在编码器和解码器之间引入了残差连接，极大地改善了梯度流（灵感来自于（He et al., 2015））。

可以看出，U-Net模型首先对输入进行下采样（即，在空间分辨率方面使输入更小），之后执行上采样。

构建Diffusion模型

import mindspore as ms
import mindspore.nn as nn
import mindspore.ops as ops

def rearrange(head, inputs):
    """
    调整输入张量的形状以适应多头注意力机制。
    head: 头的数量。
    inputs: 输入张量。
    """
    b, hc, x, y = inputs.shape  # 获取输入张量的形状
    c = hc // head  # 计算每个头的特征数量
    return inputs.reshape((b, head, c, x * y))  # 重新调整张量形状

def rsqrt(x):
    """
    计算输入的逆平方根。
    x: 输入张量。
    """
    res = ops.sqrt(x)  # 计算平方根
    return ops.inv(res)  # 计算逆操作

def randn_like(x, dtype=None):
    """
    生成与输入张量形状和数据类型相同的随机正态分布张量。
    x: 输入张量。
    dtype: 输出张量的数据类型。
    """
    if dtype is None:
        dtype = x.dtype  # 如果没有指定数据类型，则使用输入张量的数据类型
    res = ops.standard_normal(x.shape).astype(dtype)  # 生成随机正态分布张量并转换数据类型
    return res

def randn(shape, dtype=None):
    """
    生成具有指定形状和数据类型的随机正态分布张量。
    shape: 输出张量的形状。
    dtype: 输出张量的数据类型。
    """
    if dtype is None:
        dtype = ms.float32  # 如果没有指定数据类型，则使用float32
    res = ops.standard_normal(shape).astype(dtype)  # 生成随机正态分布张量并转换数据类型
    return res

def randint(low, high, size, dtype=ms.int32):
    """
    生成一个随机整数张量，其值在指定范围内。
    low: 整数范围的下限。
    high: 整数范围的上限。
    size: 输出张量的形状。
    dtype: 输出张量的数据类型。
    """
    res = ops.uniform(size, Tensor(low, dtype), Tensor(high, dtype), dtype=dtype)  # 生成均匀分布的随机数
    return res

def exists(x):
    """
    检查变量是否为None。
    x: 要检查的变量。
    """
    return x is not None  # 返回True如果x不是None

def default(val, d):
    """
    返回第一个非None的值。
    val: 第一个值。
    d: 第二个值或返回默认值的函数。
    """
    if exists(val):
        return val  # 返回第一个值如果它不是None
    return d() if callable(d) else d  # 如果第一个值为None，则返回第二个值或调用函数

def _check_dtype(d1, d2):
    """
    检查并返回支持的数据类型。
    d1: 第一个数据类型。
    d2: 第二个数据类型。
    """
    if ms.float32 in (d1, d2):
        return ms.float32  # 如果任一数据类型为float32，则返回float32
    if d1 == d2:
        return d1  # 如果两个数据类型相同，则返回该类型
    raise ValueError('dtype is not supported.')  # 抛出异常如果数据类型不支持

class Residual(nn.Cell):
    """
    实现残差连接的类。
    """
    def __init__(self, fn):
        super().__init__()  # 调用基类的构造函数
        self.fn = fn  # 保存传入的函数

    def construct(self, x, *args, **kwargs):
        """
        实现残差连接的前向传播。
        x: 输入张量。
        *args, **kwargs: 其他参数。
        """
        return self.fn(x, *args, **kwargs) + x  # 返回输入和函数输出的和

#定义上采样和下采样操作的别名
def Upsample(dim):
    return nn.Conv2dTranspose(dim, dim, 4, 2, pad_mode="pad", padding=1)

def Downsample(dim):
    return nn.Conv2d(dim, dim, 4, 2, pad_mode="pad", padding=1)

位置向量

由于神经网络的参数在时间（噪声水平）上共享，作者使用正弦位置嵌入来编码𝑡𝑡，灵感来自Transformer（Vaswani et al., 2017）。对于批处理中的每一张图像，神经网络"知道"它在哪个特定时间步长（噪声水平）上运行。

SinusoidalPositionEmbeddings模块采用(batch_size, 1)形状的张量作为输入（即批处理中几个有噪声图像的噪声水平），并将其转换为(batch_size, dim)形状的张量，其中dim是位置嵌入的尺寸。然后，我们将其添加到每个剩余块中。

import mindspore as ms
import mindspore.nn as nn
import mindspore.ops as ops
import numpy as np

class SinusoidalPositionEmbeddings(nn.Cell):
    """
    生成正弦位置嵌入的类。
    
    参数:
        dim (int): 嵌入的维度。
    """
    def __init__(self, dim):
        super().__init__()  # 调用基类的构造函数
        self.dim = dim  # 保存嵌入的维度
        half_dim = self.dim // 2  # 计算半个维度
        emb = math.log(10000) / (half_dim - 1)  # 计算对数比例
        emb = np.exp(np.arange(half_dim) * -emb)  # 计算正弦和余弦的缩放因子
        self.emb = Tensor(emb, ms.float32)  # 将缩放因子转换为MindSpore Tensor

    def construct(self, x):
        """
        生成正弦位置嵌入。
        
        参数:
            x (Tensor): 输入张量，通常包含序列的长度信息。
            
        返回:
            Tensor: 生成的位置嵌入张量。
        """
        emb = x[:, None] * self.emb[None, :]  # 扩展输入和嵌入张量以进行元素乘法
        emb = ops.concat((ops.sin(emb), ops.cos(emb)), axis=-1)  # 连接正弦和余弦值
        return emb  # 返回位置嵌入

ResNet/ConvNeXT块¶

接下来，我们定义U-Net模型的核心构建块。DDPM作者使用了一个Wide ResNet块（Zagoruyko et al., 2016），但Phil Wang决定添加ConvNeXT（Liu et al., 2022）替换ResNet，因为后者在图像领域取得了巨大成功。

在最终的U-Net架构中，可以选择其中一个或另一个，本文选择ConvNeXT块构建U-Net模型。

import mindspore as ms
import mindspore.nn as nn
import mindspore.ops as ops

def exists(x):
    """
    检查变量是否为None。
    x: 要检查的变量。
    """
    return x is not None

def c(dim, dim_out, k, padding, pad_mode='pad'):
    """
    创建一个卷积层。
    dim: 输入通道数。
    dim_out: 输出通道数。
    k: 卷积核大小。
    padding: 卷积的填充大小。
    pad_mode: 填充模式。
    """
    return nn.Conv2d(dim, dim_out, k, pad_mode=pad_mode, padding=padding)

class Block(nn.Cell):
    """
    实现一个卷积块，包含卷积、批量归一化和激活函数。
    """
    def __init__(self, dim, dim_out, groups=1):
        super().__init__()
        self.proj = c(dim, dim_out, 3, 1, pad_mode='pad')  # 创建卷积层
        self.norm = nn.GroupNorm(groups, dim_out)  # 创建批量归一化层
        self.act = nn.SiLU()  # 创建激活函数层

    def construct(self, x, scale_shift=None):
        """
        实现块的前向传播。
        
        参数:
            x (Tensor): 输入张量。
            scale_shift (tuple, optional): 缩放和平移参数。
        
        返回:
            Tensor: 输出张量。
        """
        x = self.proj(x)  # 卷积操作
        x = self.norm(x)  # 批量归一化

        if exists(scale_shift):  # 如果提供了缩放和平移参数
            scale, shift = scale_shift  # 解包参数
            x = x * (scale + 1) + shift  # 应用缩放和平移

        x = self.act(x)  # 激活函数
        return x

class ConvNextBlock(nn.Cell):
    """
    实现ConvNeXT风格的卷积块，包含深度卷积和MLP。
    """
    def __init__(self, dim, dim_out, *, time_emb_dim=None, mult=2, norm=True):
        super().__init__()
        self.mlp = (
            nn.SequentialCell(nn.GELU(), nn.Dense(time_emb_dim, dim))
            if exists(time_emb_dim)
            else None
        )  # 创建MLP，如果提供了time_emb_dim

        self.ds_conv = nn.Conv2d(dim, dim, 7, padding=3, group=dim, pad_mode="pad")  # 创建深度卷积层
        self.net = nn.SequentialCell(
            nn.GroupNorm(1, dim) if norm else nn.Identity(),  # 批量归一化或恒等操作
            nn.Conv2d(dim, dim_out * mult, 3, padding=1, pad_mode="pad"),  # 卷积
            nn.GELU(),  # GELU激活函数
            nn.GroupNorm(1, dim_out * mult),  # 批量归一化
            nn.Conv2d(dim_out * mult, dim_out, 3, padding=1, pad_mode="pad")  # 卷积
        )

        self.res_conv = nn.Conv2d(dim, dim_out, 1) if dim != dim_out else nn.Identity()  # 残差卷积

    def construct(self, x, time_emb=None):
        """
        实现ConvNeXT块的前向传播。
        
        参数:
            x (Tensor): 输入张量。
            time_emb (Tensor, optional): 时间嵌入张量。
        
        返回:
            Tensor: 输出张量。
        """
        h = self.ds_conv(x)  # 深度卷积
        if exists(self.mlp) and exists(time_emb):  # 如果MLP和时间嵌入都存在
            assert exists(time_emb), "time embedding must be passed in"
            condition = self.mlp(time_emb)  # 应用MLP
            condition = condition.expand_dims(-1).expand_dims(-1)  # 扩展维度以匹配h
            h = h + condition  # 将MLP的输出添加到深度卷积的输出

        h = self.net(h)  # 应用后续卷积和激活函数
        return h + self.res_conv(x)  # 残差连接

Attention模块

接下来，我们定义Attention模块，DDPM作者将其添加到卷积块之间。Attention是著名的Transformer架构(Vaswani et al., 2017)，在人工智能的各个领域都取得了巨大的成功，从NLP到蛋白质折叠。Phil Wang使用了两种注意力变体：一种是常规的multi-head self-attention（如Transformer中使用的），另一种是LinearAttention(Shen et al., 2018)，其时间和内存要求在序列长度上线性缩放，而不是在常规注意力中缩放。 要想对Attention机制进行深入的了解，请参照Jay Allamar的精彩的博文。

import mindspore as ms
import mindspore.nn as nn
import mindspore.ops as ops
from mindspore import Parameter, Tensor
import numpy as np

def rearrange(head, inputs):
    """
    调整输入张量的形状以适应多头注意力机制。
    head: 头的数量。
    inputs: 输入张量。
    """
    b, hc, x, y = inputs.shape
    c = hc // head
    return inputs.reshape((b, head, c, x * y))

def rsqrt(x):
    """
    计算输入的逆平方根。
    x: 输入张量。
    """
    return 1 / np.sqrt(x)

class Attention(nn.Cell):
    """
    实现多头自注意力机制的类。
    """
    def __init__(self, dim, heads=4, dim_head=32):
        super().__init__()
        self.scale = dim_head ** -0.5
        self.heads = heads
        hidden_dim = dim_head * heads

        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, pad_mode='valid', has_bias=False)
        self.to_out = nn.Conv2d(hidden_dim, dim, 1, pad_mode='valid', has_bias=True)
        self.map = ops.Map()
        self.partial = ops.Partial()

    def construct(self, x):
        b, _, h, w = x.shape
        qkv = self.to_qkv(x).chunk(3, 1)
        q, k, v = self.map(self.partial(rearrange, self.heads), qkv)

        q = q * self.scale

        sim = ops.bmm(q.swapaxes(2, 3), k)
        attn = ops.softmax(sim, axis=-1)
        out = ops.bmm(attn, v.swapaxes(2, 3))
        out = out.swapaxes(-1, -2).reshape((b, -1, h, w))

        return self.to_out(out)

class LayerNorm(nn.Cell):
    """
    实现层归一化（Layer Normalization）的类。
    """
    def __init__(self, dim):
        super().__init__()
        self.g = Parameter(initializer('ones', (1, dim, 1, 1)), name='g')

    def construct(self, x):
        eps = 1e-5
        var = x.var(1, keepdims=True)
        mean = x.mean(1, keep_dims=True)
        return (x - mean) * rsqrt((var + eps)) * self.g

class LinearAttention(nn.Cell):
    """
    实现线性复杂度的注意力机制的类。
    """
    def __init__(self, dim, heads=4, dim_head=32):
        super().__init__()
        self.scale = dim_head ** -0.5
        self.heads = heads
        hidden_dim = dim_head * heads
        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, pad_mode='valid', has_bias=False)

        self.to_out = nn.SequentialCell(
            nn.Conv2d(hidden_dim, dim, 1, pad_mode='valid', has_bias=True),
            LayerNorm(dim)
        )

        self.map = ops.Map()
        self.partial = ops.Partial()

    def construct(self, x):
        b, _, h, w = x.shape
        qkv = self.to_qkv(x).chunk(3, 1)
        q, k, v = self.map(self.partial(rearrange, self.heads), qkv)

        q = ops.softmax(q, -2)
        k = ops.softmax(k, -1)

        q = q * self.scale
        v = v / (h * w)

        context = ops.bmm(k, v.swapaxes(2, 3))
        out = ops.bmm(context.swapaxes(2, 3), q)

        out = out.reshape((b, -1, h, w))
        return self.to_out(out)

组归一化

DDPM作者将U-Net的卷积/注意层与群归一化（Wu et al., 2018）。下面，我们定义一个PreNorm类，将用于在注意层之前应用groupnorm

import mindspore as ms
import mindspore.nn as nn

class PreNorm(nn.Cell):
    """
    实现预归一化（Pre-Layer Normalization）的类。
    
    参数:
        dim (int): 输入的特征维度。
        fn (Cell): 要应用的神经网络层或函数。
    """
    def __init__(self, dim, fn):
        super(PreNorm, self).__init__()  # 调用基类的构造函数
        self.fn = fn  # 保存要应用的神经网络层或函数
        self.norm = nn.GroupNorm(1, dim)  # 创建批量归一化层，group=1表示不分组

    def construct(self, x):
        """
        实现预归一化的前向传播。
        
        参数:
            x (Tensor): 输入张量。
        
        返回:
            Tensor: 归一化后的输出张量。
        """
        x = self.norm(x)  # 应用批量归一化
        return self.fn(x)  # 应用传入的神经网络层或函数

条件U-Net

正向扩散

数据准备与处理

采样

训练过程

import mindspore as ms
from mindspore import nn

# 定义动态学习率
# 使用余弦衰减学习率调度器，设置最小学习率为1e-7，最大学习率为1e-4
# total_step为总训练步数，step_per_epoch为每个epoch的步数，decay_epoch为衰减周期
lr = nn.cosine_decay_lr(min_lr=1e-7, max_lr=1e-4, total_step=10*3750, step_per_epoch=3750, decay_epoch=10)

# 定义 Unet模型
# 假设Unet类已经在其他地方定义，且接受以下参数
# dim为特征图的维度，channels为输入通道数，dim_mults为特征图维度的乘数
unet_model = Unet(
    dim=image_size,
    channels=channels,
    dim_mults=(1, 2, 4,)
)

# 获取模型参数的名字列表
name_list = []
for (name, par) in list(unet_model.parameters_and_names()):
    name_list.append(name)

# 为模型的可训练参数设置名字
i = 0
for item in list(unet_model.trainable_params()):
    item.name = name_list[i]
    i += 1

# 定义优化器
# 使用Adam优化器，学习率由动态学习率调度器lr确定
optimizer = nn.Adam(unet_model.trainable_params(), learning_rate=lr)

# 定义动态损失尺度，用于混合精度训练
loss_scaler = ms.amp.DynamicLossScaler(65536, 2, 1000)

# 定义前向过程
def forward_fn(data, t, noise=None):
    """
    定义模型的前向传播过程。
    
    参数:
        data (Tensor): 输入数据。
        t (Tensor): 时间步长。
        noise (Tensor, optional): 噪声。
    
    返回:
        Tensor: 损失值。
    """
    loss = p_losses(unet_model, data, t, noise)  # 计算损失
    return loss

# 计算梯度
# 使用mindspore.value_and_grad计算前向过程的梯度
grad_fn = ms.value_and_grad(forward_fn, None, optimizer.parameters, has_aux=False)

# 梯度更新
def train_step(data, t, noise):
    """
    执行单步训练，包括前向传播、梯度计算和参数更新。
    
    参数:
        data (Tensor): 输入数据。
        t (Tensor): 时间步长。
        noise (Tensor): 噪声。
    
    返回:
        Tensor: 当前步的损失值。
    """
    loss, grads = grad_fn(data, t, noise)  # 计算损失和梯度
    optimizer(grads)  # 更新参数
    return loss

import time
import mindspore as ms
from mindspore import nn
import matplotlib.pyplot as plt

# 定义训练的epoch数，这里为了演示方便，设置为1
epochs = 1

# 训练循环
for epoch in range(epochs):
    # 记录当前epoch开始的时间
    begin_time = time.time()
    
    # 创建数据集的迭代器
    for step, batch in enumerate(dataset.create_tuple_iterator()):
        # 设置模型为训练模式
        unet_model.set_train()
        
        # 获取当前批次的大小
        batch_size = batch[0].shape[0]
        
        # 随机生成时间步长t
        t = ms.ops.randint(0, timesteps, (batch_size,), dtype=ms.int32)
        
        # 为当前批次生成随机噪声
        noise = ms.ops.randn_like(batch[0])
        
        # 计算当前批次的损失
        loss = train_step(batch[0], t, noise)
        
        # 每500步打印一次损失信息
        if step % 500 == 0:
            print(" epoch: ", epoch, " step: ", step, " Loss: ", loss)
    
    # 记录当前epoch结束的时间
    end_time = time.time()
    
    # 计算并打印训练时间
    times = end_time - begin_time
    print("training time:", times, "s")
    
    # 展示随机采样效果
    # 设置模型为评估模式
    unet_model.set_train(False)
    
    # 从模型中采样生成图像
    samples = sample(unet_model, image_size=image_size, batch_size=64, channels=channels)
    
    # 显示一个随机生成的图像样本
    plt.imshow(samples[-1][5].reshape(image_size, image_size, channels), cmap="gray")

# 打印训练成功的消息
print("Training Success!")

学习心得:

1. 通过学习，我深刻理解了扩散模型的核心原理，包括正向扩散过程和逆向去噪过程。这种从数据中逐步添加噪声，再通过学习逆转这一过程以生成数据的方法，展示了一种新颖的生成模型设计思路。

2. 扩散模型提供了一种不同于GANs和VAEs的生成模型范式。它基于随机过程和神经网络的结合，为生成高质量的图像提供了新的可能性。

3. 在训练扩散模型时，数据的预处理，如归一化和数据增强，对模型性能有显著影响。我学会了如何调整这些步骤以适应特定的模型和任务。

4. U-Net作为扩散模型中常用的网络结构，其编码器-解码器的设计以及残差连接的使用，对于模型的学习能力至关重要。我通过实践了解了其在图像生成中的作用。

5. 通过学习正弦位置嵌入和时间嵌入，我认识到在序列模型中引入位置信息对于模型理解数据结构的重要性。

6. 在训练过程中，使用混合精度训练可以提高训练速度并减少内存使用。我学会了如何使用MindSpore的DynamicLossScaler来实现这一点。

7. 学习扩散模型激发了我对深度学习领域的进一步兴趣，特别是对于探索新的模型结构和算法。

加油！！！