mindspore打卡第8天：DDPM的训练迭代过程和采样过程 公式显示版本
DDPM的unet模型中，训练得到的模型参数 是对于每一个时间步就生成对应一批的参数比如5个时间步对应5批不同的参数，还是说在几个时间步中的参数是迭代更新的比如一起参数随5个时间步更新还是一批？在Diffusion Models，特别是DDPM（Denoising Diffusion Probabilistic Models）中，其核心组件之一是U-Net这样的网络架构，该架构用于逐步逆转扩散过程，从而生成高质量的数据样本。关于您的问题，训练DDPM中的U-Net模型时，其参数不是针对每个时间步单独生成或更新的。
具体来说：
- **参数共享**: U-Net模型在整个扩散序列中使用的是同一套参数。这意味着不论是哪个时间步，U-Net应用的是相同的权重和偏置来处理输入。它不会为每个时间步生成一批不同的参数；相反，模型学习到如何基于当前的噪声水平（时间步）以及输入的 noisy image 来预测噪声应该如何减少，以逐步还原出清晰的样本。
- **迭代更新**: 在训练过程中，U-Net的参数是通过观察多个时间步的样本对（通常是通过扩散过程添加了不同水平噪声的数据与其对应的干净数据）并进行反向传播来迭代更新的。这意味着模型学习到了如何在不同的去噪阶段（即不同的时间步）有效操作，而这一学习过程是跨所有时间步共享知识的，不是为每个时间步独立进行的。
简而言之，DDPM中的U-Net模型训练遵循典型的深度学习训练模式，其参数是针对整个任务统一学习并在训练迭代中逐渐优化的，而非按时间步分离更新。训练得到的模型参数 是通用于采样过程的吗，采样过程中训练参数亥会更新吗？训练得到的模型参数是通用于采样过程的。一旦DDPM模型完成训练，其学习到的参数（如U-Net模型中的权重和偏置项）就被固定下来，用于指导采样过程中的去噪操作。在采样阶段，这些参数不会更新。
具体来说：
- **训练阶段：** 模型通过大量的输入数据学习如何从噪声中恢复信号，优化目标是使得模型预测的噪声尽可能接近实际添加的噪声。在这个阶段，模型的参数（θ）通过反向传播和优化算法（如Adam）不断调整，以最小化损失函数。
- **采样（生成）阶段：** 训练完成后，模型参数被视为最优解并固定。采样时，我们利用这些固定的参数逐步从噪声中构建出清晰的图像。这意味着对于任意新的输入或采样请求，模型都将使用相同的参数执行去噪过程，而不再需要额外的学习或参数更新。
因此，训练阶段和采样阶段是分开的：前者是学习阶段，后者是应用阶段。采样时模型扮演的是一个生成器角色，利用学到的知识生成新样本，而不是继续学习或调整参数。讲讲DDPM中unet模型 在训练步骤和采样步骤中是怎么一个作用和流程？在扩散模型（DDPM，Denoising Diffusion Probabilistic Models）中，U-Net模型通常扮演着去噪器的角色，它被用来学习如何从噪声中逐渐恢复出清晰的图像数据。下面分别说明U-Net模型在DDPM的训练步骤和采样步骤中的作用和流程：
### 训练步骤：
1. **前向过程（扩散过程）：** 首先，给定一张干净的图像 \(x_0\)，通过一系列逐步增加噪声的步骤生成一系列中间噪声图像 \(x_1, x_2, ..., x_T\)，直到 \(x_T\) 接近纯噪声。这个过程可以视为一个马尔科夫链，每一步添加的噪声量由一个预定的方差决定。
2. **学习去噪：** 在训练阶段，模型不是直接学习从噪声到图像的映射，而是学习如何预测在给定时间步 \(t\) 下，原始图像 \(x_0\) 添加了噪声后变成 \(x_t\) 时的噪声分布。具体来说，U-Net模型被训练来预测 \(\epsilon_\theta(x_t, t)\)，即给定时间步 \(t\) 和当前噪声图像 \(x_t\)，模型尝试估计原本添加到图像上的噪声。
3. **损失函数：** 训练的目标是让U-Net模型预测的噪声尽可能接近实际添加的噪声。损失函数通常基于均方误差(MSE)计算预测的噪声和实际噪声之间的差异。这样，模型通过反向传播和梯度下降来更新其参数，以优化去噪能力。
### 采样步骤：
1. **初始化与逆过程：** 采样开始时，首先从一个纯粹的噪声分布（如高斯噪声）中生成一个图像 \(x_T\)。
2. **逐步去噪：** 接下来，通过迭代的方式，从时间步 \(T\) 开始逐步向前（向时间步 \(t=0\)），每一步都要使用U-Net模型去噪。在每一步中，模型会根据当前的噪声图像 \(x_t\) 和时间步 \(t\)，预测如何去掉一部分噪声，从而得到一个较为清晰的图像 \(x_{t-1}\)。
3. **采样策略：** 实际应用中，通常会采用一些策略来更好地控制采样过程，比如使用分数时间步、层级采样或者预测噪声残差等技巧。每一步中，模型预测的是如何从 \(x_t\) 转换到 \(x_{t-1}\)，这需要结合时间步 \(t\) 的特定噪声水平进行调整。
4. **终止于清晰图像：** 经过一系列这样的去噪步骤，最终会得到一个接近原始数据分布的图像 \(x_0\)，即完成了从噪声到清晰图像的生成过程。
综上所述，在DDPM中，U-Net模型通过在训练阶段学习如何预测和去除噪声，而在采样阶段则利用这一学习成果，逐步逆向去噪，从纯粹的噪声中恢复出高质量的图像数据。在DDPM代码中，看到train_step这里的epoch=1， t是一个10维的tensor进行输入，并没有看到tain_step含有sample迭代的代码。。在训练阶段unet作为model作用是啥，他实现了什么功能？？而当在sample的阶段时候输入的是大T，这个体现了时间步的迭代。训练阶段的模型体现不出迭代的过程，为什么？在DDPM的训练过程中，`train_step`函数执行的是单次迭代的训练逻辑，其中`epoch=1`指示的是当前正在进行第一个训练周期。尽管在这个描述中没有直接提及样本迭代的代码段，但训练迭代实际上通过数据加载器的循环在更高级别实现。也就是说，对于每个epoch，数据加载器会顺序或随机地提供一批批的数据（batch），而`train_step`针对每个这样的批次运行一次，从而在整体上遍历整个训练数据集，实现了样本的迭代训练。
U-Net作为DDPM中的核心模型，在训练阶段扮演着关键角色，其主要功能包括：
1. **学习逆向过程**：U-Net试图学习如何从加噪图像逆向恢复原始图像。在每次迭代中，给定一个扩散时间`t`（在这个场景中，`t`是10维的张量，意味着同时处理10个不同的扩散时间点），U-Net尝试预测如何去噪，即学习噪声分布或者直接预测如何减少当前图像中的噪声，使其更接近真实图像。
2. **捕获多尺度特征**：由于U-Net的结构特性，它能够有效地捕获和利用不同尺度的特征，这对于从高度噪声化的图像恢复细节非常重要。这种能力使得U-Net在逐步去噪的过程中能有效提升图像质量。
至于“训练阶段的模型体现不出迭代的过程”，这里可能存在一些理解上的混淆。实际上，训练过程本身就是一个迭代过程，它通过不断调整模型参数来最小化损失函数，这本身就是一种迭代优化。每次迭代（`train_step`的执行）都会基于当前的模型参数和训练样本计算梯度，并据此更新参数。因此，即使在单次`train_step`调用中没有直接看到时间步的逐级递减，整个训练流程（跨越多个epoch和批次）确实体现了迭代的性质，每一次迭代都是对模型的一次微调，累积起来使得模型逐渐学习到如何逆向去噪。
而在采样阶段，输入大`T`并逐步减小到0，是通过一系列连续的U-Net预测来体现时间步的迭代，这是在训练阶段学到的能力的实际应用，通过模拟反向扩散过程逐步从噪声中生成清晰图像。所以，尽管训练和采样阶段的迭代表现形式不同，它们共同构成了DDPM工作流程的核心。
在DDPM的训练过程中，预测的噪声确实体现了迭代的时间步特征，这一点在训练阶段的数学公式中有所体现。为了更好地理解这一点，我们可以回顾一下DDPM训练过程中的核心公式。
### 扩散过程 (Forward Process)
扩散过程是将数据逐渐添加高斯噪声的过程。设$x_t$表示在时间$t$的数据状态，那么扩散过程可以用下面的公式表示：
\[
x_t = \sqrt{\alpha_t} x_{t-1} + \sqrt{1 - \alpha_t} \epsilon \quad \text{with } \epsilon \sim \mathcal{N}(0, I)
\]
其中，$\alpha_t$是一个介于0和1之间的衰减系数，随时间$t$变化。$\epsilon$是从标准正态分布$\mathcal{N}(0, I)$中抽样的随机噪声。
### 去噪过程 (Reverse Process)
去噪过程是训练模型的目标，即学习一个条件概率$p_\theta(x_{t-1}|x_t)$，它描述了如何从$x_t$中恢复$x_{t-1}$。在DDPM中，我们假设$p_\theta(x_{t-1}|x_t)$可以被近似为高斯分布，其均值由U-Net模型预测的噪声$\epsilon_\theta(x_t, t)$决定：
\[
p_\theta(x_{t-1}|x_t) = \mathcal{N}\left(x_{t-1}; \mu_\theta(x_t, t), \sigma^2_t I\right)
\]
其中，$\mu_\theta(x_t, t)$是U-Net模型输出，表示预测的去噪方向；$\sigma^2_t$是一个固定的时间步依赖的方差。
### 训练目标
训练的目标是最大化数据在所有时间步的后验概率的对数似然估计。然而，由于直接计算所有时间步的概率非常困难，DDPM采用了简化的目标，即最小化噪声预测的均方误差：
\[
L_\text{simple}(\theta) = \mathbb{E}_{t, x_0, \epsilon}\left[\|\epsilon - \epsilon_\theta(\sqrt{\bar{\alpha}_t} x_0 + \sqrt{1-\bar{\alpha}_t}\epsilon, t)\|^2\right]
\]
其中，$\bar{\alpha}_t = \prod_{s=1}^{t}\alpha_s$表示从$t=0$到$t$的累积衰减系数，$\epsilon$是从标准正态分布中抽样的随机噪声，$x_0$是原始的清晰图像，$t$是从$1$到$T$随机选择的时间步。
在训练的每一个迭代中，模型都会从数据集中抽取一批图像$x_0$，然后随机选择一个时间步$t$，并根据上述公式构造带噪声的图像$x_t$。模型的输出$\epsilon_\theta(x_t, t)$是对噪声$\epsilon$的预测。通过最小化预测噪声与实际噪声之间的差距，模型逐步学习了如何根据当前的时间步$t$和当前的噪声图像$x_t$预测并去除噪声。
### 时间步的迭代特征
在训练过程中，虽然模型在每个时间步$t$上只做一次预测，但通过不断地迭代训练，模型逐渐学会了在所有可能的时间步上如何有效地预测噪声。换句话说，模型在训练中不断接触到不同时间步下的数据和噪声，从而能够学会如何在任何给定时间步上预测和去除噪声，这正是时间步迭代特征的体现。
在采样阶段，模型使用在训练阶段学到的技能，通过从噪声开始逐步预测并去除噪声，最终恢复出清晰的图像，这一过程本质上是在时间步上的逆向迭代，从$T$到$0$，每一步都利用了模型在相应时间步上预测噪声的能力。这不仅验证了模型在训练阶段确实学到了时间步的迭代特征，而且展示了这些特征在实际生成过程中的应用。
上图展示了DDPM中去噪过程的可视化效果，从几乎完全随机的噪声图像逐步恢复到清晰的原始图像，这一过程体现了时间步的迭代特征。可以看到，随着去噪过程的推进，图像细节逐渐显现，最终呈现出清晰可辨的形态，直观地展现了从高斯噪声中恢复图像信息的过程，以及模型在不同时间步上预测和去除噪声的能力。
在去噪扩散概率模型（DDPM）中，简化的目标函数通常被设计成直接预测噪声项，而不是逐步预测每一时间步的去噪结果。这是因为直接预测噪声比预测每一时间步的图像更稳定且更容易优化。在DDPM中，模型的目标是在给定时间步\(t\)的噪声图像\(x_t\)的情况下，预测添加的噪声\(\epsilon\)。
让我们回顾一下DDPM的简化目标函数，即最小化噪声预测的均方误差（MSE）：
\[ L_{\text{simple}}(\theta) = \mathbb{E}_{t, x_0, \epsilon}\left[\left\|\epsilon - \epsilon_\theta\left(\sqrt{\bar{\alpha}_t} x_0 + \sqrt{1-\bar{\alpha}_t}\epsilon, t\right)\right\|^2\right] \]
其中，
- \(x_0\)是原始的清晰图像，
- \(\epsilon\)是从标准正态分布中抽样的随机噪声，
- \(t\)是从1到T随机选择的时间步，
- \(\bar{\alpha}_t\)是累积衰减系数，\(\bar{\alpha}_t = \prod_{s=1}^{t}\alpha_s\)。
### 推导过程
简化的目标函数实际上是基于后验概率的近似和变分下界（ELBO）推导而来。在DDPM中，我们关注的是最大化数据\(x_0\)的对数似然，即\(\log p_\theta(x_0)\)。然而，直接优化这个目标是不可行的，因为涉及到复杂的积分和高维分布。因此，我们使用变分方法来逼近这个目标。
在训练阶段，我们假设模型\(p_\theta(x_{t-1}|x_t)\)可以被近似为高斯分布，其中的均值由U-Net模型预测的噪声\(\epsilon_\theta(x_t, t)\)确定。具体来说，我们可以将后验概率\(p_\theta(x_{t-1}|x_t)\)写作：
\[ p_\theta(x_{t-1}|x_t) = \mathcal{N}\left(x_{t-1}; \mu_\theta(x_t, t), \sigma^2_t I\right) \]
其中，\(\mu_\theta(x_t, t)\)是U-Net模型输出，表示预测的去噪方向；\(\sigma^2_t\)是一个固定的时间步依赖的方差。
由于直接优化后验概率的对数似然是困难的，我们转而优化一个更简单的替代目标，即预测噪声的MSE。这个简化的目标函数可以被证明是后验概率的对数似然的一个下界，因此优化它可以间接促进对数似然的提高。
### 直接预测噪声的优势
直接预测噪声而不是每一时间步的去噪图像有以下几个优势：
1. **稳定性**：预测噪声通常比预测图像本身更加稳定，因为它避免了累积误差的问题。
2. **计算效率**：预测噪声的模型可以更快收敛，因为噪声分布相对简单。
3. **泛化能力**：噪声预测模型往往具有更好的泛化能力，能够在未见过的数据上表现更好。
因此，DDPM通过简化目标函数，专注于预测噪声，实现了高效的训练和高质量的生成结果。
在深度学习和生成模型领域，最大化数据\(x_0\)的对数似然\(\log p_\theta(x_0)\)意味着我们试图找到一组参数\(\theta\)，使得模型\(p_\theta\)能够尽可能准确地生成与\(x_0\)相似的数据。在图像生成任务中，\(x_0\)代表原始的、未受噪声影响的图像，而\(\log p_\theta(x_0)\)则衡量了模型生成类似\(x_0\)的图像的概率。
### 理解\(\log p_\theta(x_0)\)
对于一张图像\(x_0\)，其像素值可以视为多维随机变量的观测值，这些随机变量遵循某种概率分布。在高维空间中，直接计算\(p_\theta(x_0)\)是非常困难的，因为这涉及到对所有可能的像素组合进行积分或求和，而图像的像素数目通常很大。例如，一张\(64 \times 64\)的灰度图像就有\(64 \times 64 = 4096\)个像素，每个像素可以取从0到255共256种不同的值，这意味着总的可能性数量级非常巨大。
取对数\(\log p_\theta(x_0)\)有几个关键的好处：
1. **数值稳定性**：直接处理概率\(p_\theta(x_0)\)可能会导致非常小的数值，而在计算机中存储和计算这些极小数值时容易产生数值不稳定问题。对数可以将乘法转换为加法，从而避免了这种问题。
2. **梯度计算**：在优化过程中，我们通常需要计算损失函数的梯度。对数似然函数的梯度通常更容易计算，尤其是在处理复杂的概率模型时。
### 数据举例
假设我们有一张尺寸为\(32 \times 32\)像素的RGB图像，这意味着它有\(32 \times 32 \times 3 = 3072\)个像素。每个像素可以有256种可能的值（如果我们将颜色范围限定在0到255之间）。\(x_0\)可以看作是这个3072维空间中的一个点，其每个坐标值对应一个像素的红、绿、蓝三通道的强度值。
现在，假设我们有一个深度生成模型，比如变分自编码器（VAE）或生成对抗网络（GAN），它的目标是学习一个参数化的概率分布\(p_\theta\)，这个分布能够生成看起来与\(x_0\)相似的新图像。如果我们使用\(p_\theta\)来生成一幅图像，并且这幅图像与\(x_0\)非常接近，那么\(p_\theta(x_0)\)就会相对较大。然而，由于\(p_\theta(x_0)\)在实际中是一个非常小的数，我们通常会关注其对数值\(\log p_\theta(x_0)\)，这样可以避免数值溢出或下溢问题。
在训练过程中，我们的目标就是调整参数\(\theta\)，以使\(\log p_\theta(x_0)\)尽可能大，这意味着模型\(p_\theta\)生成与\(x_0\)相似图像的概率更高。但直接优化\(\log p_\theta(x_0)\)往往是不可行的，特别是在复杂的模型中，因此我们通常会采用各种技巧和近似方法来间接优化这个目标，比如使用变分下界（ELBO）或通过预测噪声等方式。这些技巧使得我们可以有效地训练模型，同时保持对原始目标的追求——即生成与训练数据相似的新样本。
其中，\(\mu_\theta(x_t, t)\)是U-Net模型输出，表示预测的去噪方向；\(\sigma^2_t\)是一个固定的时间步依赖的方差。 在去噪扩散概率模型（DDPM）中，简化的目标函数通常被设计成直接预测噪声项，而不是逐步预测每一时间步的去噪结果。 这个详细解析一下
去噪扩散概率模型（DDPM，Denoising Diffusion Probabilistic Models）是一种生成模型，它通过向数据中添加高斯噪声，然后训练一个模型来逐步去除这些噪声，从而恢复数据。这一过程涉及两个阶段：前向扩散过程和后向去噪过程。
### 前向扩散过程
在前向扩散过程中，模型会将数据逐渐转化为噪声。具体来说，从干净的数据\(x_0\)开始，模型会在多个时间步\(t=1,2,...,T\)中，逐步向数据中添加高斯噪声，得到噪声化数据\(x_t\)。这一过程可以用以下公式表示：
\[q(x_t|x_0) = \mathcal{N}(x_t; \sqrt{\bar{\alpha}_t}x_0, (1-\bar{\alpha}_t)\mathbf{I})\]
其中，\(\bar{\alpha}_t = \prod_{s=1}^{t}\alpha_s\)，而\(\alpha_s\)和\(1-\alpha_s\)分别表示保留信号和添加噪声的比例。在这个过程中，\(x_t\)逐渐变成噪声，直到\(t=T\)时，\(x_T\)几乎完全由噪声组成。
### 后向去噪过程
后向去噪过程的目标是通过逆向传播，从噪声数据\(x_T\)逐步恢复出干净的数据\(x_0\)。这一过程是通过训练一个神经网络模型，如U-Net，来预测每个时间步的噪声项，而不是直接预测去噪后的结果。U-Net模型的输出\(\mu_\theta(x_t, t)\)代表了在时间步\(t\)下，给定噪声数据\(x_t\)，模型预测的去噪方向，即噪声项\(\epsilon\)的估计。
DDPM的核心思想是，通过预测并移除噪声，可以逐步逼近数据的真实分布。在每个时间步\(t\)，模型预测的噪声项被用于生成下一个更清晰的图像版本\(x_{t-1}\)。这一过程可以表示为：
\[p_\theta(x_{t-1}|x_t) = \mathcal{N}(x_{t-1}; \mu_\theta(x_t, t), \beta_t\mathbf{I})\]
其中，\(\mu_\theta(x_t, t)\)是U-Net模型的输出，\(\beta_t\)是时间步\(t\)的方差。模型的训练目标是使\(\mu_\theta(x_t, t)\)尽可能接近真实的噪声项\(\epsilon\)。
### 目标函数
在训练过程中，DDPM的目标函数通常设计成最小化噪声预测的均方误差（MSE）。具体来说，模型尝试最小化以下损失函数：
\[\mathbb{E}_{t,x_0,\epsilon}[(\epsilon - \mu_\theta(\sqrt{\bar{\alpha}_t}x_0 + \sqrt{1-\bar{\alpha}_t}\epsilon, t))^2]\]
这里，\(\epsilon\)是在每个时间步\(t\)上独立采样得到的噪声，而\(\mu_\theta\)是U-Net模型在时间步\(t\)下的输出，它试图预测这个噪声项。通过最小化上述损失函数，模型可以学习到如何在每个时间步上预测和去除噪声，最终达到从噪声中恢复出原始数据的效果。
简而言之，DDPM通过将复杂的数据生成过程分解为一系列去噪步骤，使得模型能够逐步从噪声中恢复出原始数据，而这一过程的关键在于训练模型预测每个时间步上的噪声项，而非直接预测去噪后的数据本身。上图形象地展示了去噪扩散概率模型（DDPM）从噪声图像中逐步恢复出干净图像的过程。从左至右，我们可以看到图像从高度失真的状态逐渐变得清晰，最终恢复出原始的细节。这一过程体现了DDPM模型通过逐步去除噪声，成功重建图像的能力。在实际应用中，这一过程是由模型在多个时间步上预测并移除噪声项实现的，正如之前解析中所描述的那样。
在训练过程中，DDPM的目标函数通常设计成最小化噪声预测的均方误差（MSE）。具体来说，模型尝试最小化以下损失函数：
\[\mathbb{E}_{t,x_0,\epsilon}[(\epsilon - \mu_\theta(\sqrt{\bar{\alpha}_t}x_0 + \sqrt{1-\bar{\alpha}_t}\epsilon, t))^2]\]。。，\(\mu_\theta(x_t, t)\)是U-Net模型输出，表示预测的去噪方向怎么理解去噪方向
在DDPM中，术语“去噪方向”是指模型预测的噪声项的方向，这实质上就是模型对于当前时刻\(t\)的输入\(x_t\)认为应该减去的噪声部分。让我们更深入地理解这一点：
### \(x_t\)的构成
在任何时间步\(t\)，\(x_t\)都可以被看作是原始图像\(x_0\)与在该时间步上叠加的噪声\(\epsilon\)的线性组合。具体来说，
\[
x_t = \sqrt{\bar{\alpha}_t}x_0 + \sqrt{1-\bar{\alpha}_t}\epsilon
\]
这里的\(\bar{\alpha}_t\)是累积衰减因子，决定了\(x_0\)在\(x_t\)中的权重，而\(\sqrt{1-\bar{\alpha}_t}\epsilon\)则表示在时间步\(t\)上叠加的噪声成分。
### 模型的预测
当模型\(\mu_\theta(x_t, t)\)接收\(x_t\)作为输入时，其目标是预测\(\epsilon\)，即被添加到\(x_0\)以产生\(x_t\)的噪声。由于\(\epsilon\)本质上是\(x_t\)和\(x_0\)之间的差异，因此预测\(\epsilon\)可以被视为确定了从\(x_t\)“回归”到\(x_0\)的方向——即去噪方向。
### 如何使用预测的去噪方向
一旦模型预测出了\(\epsilon\)，我们就可以用它来计算\(x_{t-1}\)，即离\(x_0\)更近一步的图像。具体计算方式涉及到将预测的\(\epsilon\)从\(x_t\)中减去，并适当地缩放以得到\(x_{t-1}\)。这一过程可以写作：
\[
x_{t-1} = \frac{1}{\sqrt{\alpha_t}}(x_t - \frac{\sqrt{1-\bar{\alpha}_t}}{\sqrt{1-\alpha_t}}\mu_\theta(x_t, t))
\]
这里\(\alpha_t\)是单步衰减因子，\(\frac{\sqrt{1-\bar{\alpha}_t}}{\sqrt{1-\alpha_t}}\mu_\theta(x_t, t)\)就是我们说的去噪方向，它指明了如何修正\(x_t\)以接近\(x_0\)。
### 目标函数的作用
目标函数
\[
\mathbb{E}_{t,x_0,\epsilon}[(\epsilon - \mu_\theta(\sqrt{\bar{\alpha}_t}x_0 + \sqrt{1-\bar{\alpha}_t}\epsilon, t))^2]
\]
实际上是在度量模型预测的去噪方向（即\(\mu_\theta(x_t, t)\)）与真实的噪声项\(\epsilon\)之间的差距。通过最小化这个差距，模型学习到如何准确预测噪声，从而在每个时间步上有效地去除噪声，最终恢复出原始的图像\(x_0\)。
综上所述，“去噪方向”是指模型预测的噪声\(\epsilon\)，这一预测指导着如何从噪声化的图像中逐步恢复出原始图像，通过逐步减少噪声，模型能够重建出原始数据。
在训练过程中，DDPM的目标函数通常设计成最小化噪声预测的均方误差（MSE）。具体来说，模型尝试最小化以下损失函数：
\[\mathbb{E}_{t,x_0,\epsilon}[(\epsilon - \mu_\theta(\sqrt{\bar{\alpha}_t}x_0 + \sqrt{1-\bar{\alpha}_t}\epsilon, t))^2]\]。。，\(\mu_\theta(x_t, t)\)是U-Net模型输出，表示预测的去噪方向。在训练过程中，虽然模型在每个时间步$t$上只做一次预测，但通过不断地迭代训练，模型逐渐学会了在所有可能的时间步上如何有效地预测噪声。换句话说，模型在训练中不断接触到不同时间步下的数据和噪声，从而能够学会如何在任何给定时间步上预测和去除噪声，这正是时间步迭代特征的体现。
您提到的正是DDPM（Denoising Diffusion Probabilistic Models）训练机制的核心思想。在DDPM中，模型的目标是学习一个去噪过程，即逆向扩散过程，该过程能将高斯噪声逐步转换回原始数据分布。这一过程的关键在于模型学习如何预测并去除在不同时间步下添加的噪声，进而重构出原始数据。
### 时间步迭代特征
在训练阶段，模型在每个时间步\(t\)上都尝试预测对应的噪声\(\epsilon\)。这个预测是基于在时间步\(t\)上的数据点\(x_t\)进行的，其中\(x_t\)是通过在原始数据\(x_0\)上添加高斯噪声得到的。具体而言，\(x_t\)由下面的公式给出：
\[
x_t = \sqrt{\bar{\alpha}_t}x_0 + \sqrt{1-\bar{\alpha}_t}\epsilon
\]
其中\(\bar{\alpha}_t\)是一个随时间步变化的累积衰减系数，\(\epsilon\)是标准正态分布的随机噪声。
### 去噪方向与预测
模型\(\mu_\theta(x_t, t)\)的任务是预测\(\epsilon\)，即在给定时间步\(t\)上所添加的噪声。这个预测结果代表了从\(x_t\)回到\(x_0\)的方向和程度，也就是所谓的“去噪方向”。通过最小化预测噪声\(\mu_\theta(x_t, t)\)与真实噪声\(\epsilon\)之间的MSE（均方误差），模型学习到了在不同时间步下有效预测和去除噪声的能力。
### 训练过程与时间步迭代
训练过程涉及在不同的时间步上重复这一预测过程。模型在训练集上迭代更新其参数\(\theta\)，以优化上述定义的损失函数。在每个训练epoch中，模型会经历所有或大部分可能的时间步，这意味着它必须学会在时间序列的各个点上进行有效的去噪预测。随着时间的推移，模型逐渐优化了其参数，提高了在任意给定时间步上预测和去除噪声的准确性，这正是通过时间步迭代特征实现的。
### 结论
DDPM的训练过程通过让模型在一系列时间步上迭代预测和去除噪声，使模型能够学习到一种通用的去噪策略，这种策略适用于从高斯噪声中恢复出原始数据分布。这一机制确保了模型不仅在训练数据上表现良好，而且在处理未见过的数据时也能有效地执行去噪任务，从而生成高质量的样本。
实际上，在DDPM中，均方误差（MSE）的计算是在单个时间步上进行的，但整体的损失函数确实是所有可能时间步上MSE的期望值或者说平均值。这是因为模型需要在每一个可能的时间步\(t\)上都学会有效地预测噪声\(\epsilon\)。
具体地，对于一个特定的时间步\(t\)，模型\(\mu_\theta(x_t, t)\)的输出和真实噪声\(\epsilon\)之间的MSE可以写为：
\[
L_t = \mathbb{E}[(\epsilon - \mu_\theta(x_t, t))^2]
\]
这里的\(L_t\)表示在时间步\(t\)上的损失。然而，在整个训练过程中，模型的总损失函数是所有时间步上损失的平均，可以表示为：
\[
L = \mathbb{E}_t[L_t] = \mathbb{E}_t\left[\mathbb{E}_{x_0,\epsilon}\left[(\epsilon - \mu_\theta(x_t, t))^2\right]\right]
\]
这里的\(\mathbb{E}_t\)表示对所有时间步\(t\)取平均，而\(\mathbb{E}_{x_0,\epsilon}\)表示对数据\(x_0\)和噪声\(\epsilon\)的分布取平均。
因此，尽管模型在每次训练迭代中只在当前随机选择的时间步上计算MSE，但整体的训练目标是通过迭代优化来最小化所有时间步上的平均MSE。这确保了模型能够学习到一个通用的噪声预测策略，而不是仅针对某几个特定时间步的局部最优解。最终，模型将能够有效地在任意给定的时间步上预测并去除噪声，从而完成从高斯噪声中重建原始数据的任务。
```python
# %%capture captured_output
# # 实验环境已经预装了mindspore==2.2.14，如需更换mindspore版本，可更改下面mindspore的版本号
# !pip uninstall mindspore -y
# !pip install -i https://pypi.mirrors.ustc.edu.cn/simple mindspore==2.2.14
```
```python
# # 查看当前 mindspore 版本
# !pip show mindspore
```
```python
import math
from functools import partial
%matplotlib inline
import matplotlib.pyplot as plt
from tqdm.auto import tqdm
import numpy as np
from multiprocessing import cpu_count
from download import download
import mindspore as ms
import mindspore.nn as nn
import mindspore.ops as ops
from mindspore import Tensor, Parameter
from mindspore import dtype as mstype
from mindspore.dataset.vision import Resize, Inter, CenterCrop, ToTensor, RandomHorizontalFlip, ToPIL
from mindspore.common.initializer import initializer
from mindspore.amp import DynamicLossScaler
ms.set_seed(0)
```
## 构建Diffusion模型
下面，我们逐步构建Diffusion模型。
首先，我们定义了一些帮助函数和类，这些函数和类将在实现神经网络时使用。
```python
def rearrange(head, inputs):
b, hc, x, y = inputs.shape
c = hc // head
return inputs.reshape((b, head, c, x * y))
def rsqrt(x):
res = ops.sqrt(x)
return ops.inv(res)
def randn_like(x, dtype=None):
if dtype is None:
dtype = x.dtype
res = ops.standard_normal(x.shape).astype(dtype)
return res
def randn(shape, dtype=None):
if dtype is None:
dtype = ms.float32
res = ops.standard_normal(shape).astype(dtype)
return res
def randint(low, high, size, dtype=ms.int32):
res = ops.uniform(size, Tensor(low, dtype), Tensor(high, dtype), dtype=dtype)
return res
def exists(x):
return x is not None
def default(val, d):
if exists(val):
return val
return d() if callable(d) else d
def _check_dtype(d1, d2):
if ms.float32 in (d1, d2):
return ms.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):
return self.fn(x, *args, **kwargs) + x
```
我们还定义了上采样和下采样操作的别名。
```python
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)
```
### 位置向量
由于神经网络的参数在时间（噪声水平）上共享，作者使用正弦位置嵌入来编码$t$，灵感来自Transformer（[Vaswani et al., 2017](https://arxiv.org/abs/1706.03762)）。对于批处理中的每一张图像，神经网络"知道"它在哪个特定时间步长（噪声水平）上运行。
`SinusoidalPositionEmbeddings`模块采用`(batch_size, 1)`形状的张量作为输入（即批处理中几个有噪声图像的噪声水平），并将其转换为`(batch_size, dim)`形状的张量，其中`dim`是位置嵌入的尺寸。然后，我们将其添加到每个剩余块中。
```python
class SinusoidalPositionEmbeddings(nn.Cell):
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)
def construct(self, x):
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](https://arxiv.org/abs/1605.07146)），但Phil Wang决定添加ConvNeXT（[Liu et al., 2022](https://arxiv.org/abs/2201.03545)）替换ResNet，因为后者在图像领域取得了巨大成功。
在最终的U-Net架构中，可以选择其中一个或另一个，本文选择ConvNeXT块构建U-Net模型。
```python
class Block(nn.Cell):
def __init__(self, dim, dim_out, groups=1):
super().__init__()
self.proj = nn.Conv2d(dim, dim_out, 3, pad_mode="pad", padding=1)
self.proj = c(dim, dim_out, 3, padding=1, pad_mode='pad')
self.norm = nn.GroupNorm(groups, dim_out)
self.act = nn.SiLU()
def construct(self, x, scale_shift=None):
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):
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
)
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(),
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):
h = self.ds_conv(x)
if exists(self.mlp) and exists(time_emb):
assert exists(time_emb), "time embedding must be passed in"
condition = self.mlp(time_emb)
condition = condition.expand_dims(-1).expand_dims(-1)
h = h + condition
h = self.net(h)
return h + self.res_conv(x)
```
### Attention模块
接下来，我们定义Attention模块，DDPM作者将其添加到卷积块之间。Attention是著名的Transformer架构([Vaswani et al., 2017](https://arxiv.org/abs/1706.03762))，在人工智能的各个领域都取得了巨大的成功，从NLP到[蛋白质折叠](https://www.deepmind.com/blog/alphafold-a-solution-to-a-50-year-old-grand-challenge-in-biology)。Phil Wang使用了两种注意力变体：一种是常规的multi-head self-attention（如Transformer中使用的），另一种是[LinearAttention](https://github.com/lucidrains/linear-attention-transformer)([Shen et al., 2018](https://arxiv.org/abs/1812.01243))，其时间和内存要求在序列长度上线性缩放，而不是在常规注意力中缩放。
要想对Attention机制进行深入的了解，请参照Jay Allamar的[精彩的博文](https://jalammar.github.io/illustrated-transformer/)。
```python
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
# 'b h d i, b h d j -> b h i j'
sim = ops.bmm(q.swapaxes(2, 3), k)
attn = ops.softmax(sim, axis=-1)
# 'b h i j, b h d j -> b h i d'
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):
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)
# 'b h d n, b h e n -> b h d e'
context = ops.bmm(k, v.swapaxes(2, 3))
# 'b h d e, b h d n -> b h e n'
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](https://arxiv.org/abs/1803.08494)）。下面，我们定义一个`PreNorm`类，将用于在注意层之前应用groupnorm。
```python
class PreNorm(nn.Cell):
def __init__(self, dim, fn):
super().__init__()
self.fn = fn
self.norm = nn.GroupNorm(1, dim)
def construct(self, x):
x = self.norm(x)
return self.fn(x)
```
### 条件U-Net
我们已经定义了所有的构建块（位置嵌入、ResNet/ConvNeXT块、Attention和组归一化），现在需要定义整个神经网络了。请记住，网络 $\mathbf{\epsilon}_\theta(\mathbf{x}_t, t)$ 的工作是接收一批噪声图像+噪声水平，并输出添加到输入中的噪声。
更具体的：
网络获取了一批`(batch_size, num_channels, height, width)`形状的噪声图像和一批`(batch_size, 1)`形状的噪音水平作为输入，并返回`(batch_size, num_channels, height, width)`形状的张量。
网络构建过程如下：
- 首先，将卷积层应用于噪声图像批上，并计算噪声水平的位置
- 接下来，应用一系列下采样级。每个下采样阶段由2个ResNet/ConvNeXT块 + groupnorm + attention + 残差连接 + 一个下采样操作组成
- 在网络的中间，再次应用ResNet或ConvNeXT块，并与attention交织
- 接下来，应用一系列上采样级。每个上采样级由2个ResNet/ConvNeXT块+ groupnorm + attention + 残差连接 + 一个上采样操作组成
- 最后，应用ResNet/ConvNeXT块，然后应用卷积层
最终，神经网络将层堆叠起来，就像它们是乐高积木一样（但重要的是[了解它们是如何工作的](http://karpathy.github.io/2019/04/25/recipe/)）。
```python
class Unet(nn.Cell):
def __init__(
self,
dim,
init_dim=None,
out_dim=None,
dim_mults=(1, 2, 4, 8),
channels=3,
with_time_emb=True,
convnext_mult=2,
):
super().__init__()
self.channels = channels
init_dim = default(init_dim, dim // 3 * 2)
self.init_conv = nn.Conv2d(channels, init_dim, 7, padding=3, pad_mode="pad", has_bias=True)
dims = [init_dim, *map(lambda m: dim * m, dim_mults)]
in_out = list(zip(dims[:-1], dims[1:]))
block_klass = partial(ConvNextBlock, mult=convnext_mult)
if with_time_emb:
time_dim = dim * 4
self.time_mlp = nn.SequentialCell(
SinusoidalPositionEmbeddings(dim),
nn.Dense(dim, time_dim),
nn.GELU(),
nn.Dense(time_dim, time_dim),
)
else:
time_dim = None
self.time_mlp = None
self.downs = nn.CellList([])
self.ups = nn.CellList([])
num_resolutions = len(in_out)
for ind, (dim_in, dim_out) in enumerate(in_out):
is_last = ind >= (num_resolutions - 1)
self.downs.append(
nn.CellList(
[
block_klass(dim_in, dim_out, time_emb_dim=time_dim),
block_klass(dim_out, dim_out, time_emb_dim=time_dim),
Residual(PreNorm(dim_out, LinearAttention(dim_out))),
Downsample(dim_out) if not is_last else nn.Identity(),
]
)
)
mid_dim = dims[-1]
self.mid_block1 = block_klass(mid_dim, mid_dim, time_emb_dim=time_dim)
self.mid_attn = Residual(PreNorm(mid_dim, Attention(mid_dim)))
self.mid_block2 = block_klass(mid_dim, mid_dim, time_emb_dim=time_dim)
for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])):
is_last = ind >= (num_resolutions - 1)
self.ups.append(
nn.CellList(
[
block_klass(dim_out * 2, dim_in, time_emb_dim=time_dim),
block_klass(dim_in, dim_in, time_emb_dim=time_dim),
Residual(PreNorm(dim_in, LinearAttention(dim_in))),
Upsample(dim_in) if not is_last else nn.Identity(),
]
)
)
out_dim = default(out_dim, channels)
self.final_conv = nn.SequentialCell(
block_klass(dim, dim), nn.Conv2d(dim, out_dim, 1)
)
def construct(self, x, time):
x = self.init_conv(x)
t = self.time_mlp(time) if exists(self.time_mlp) else None
h = []
for block1, block2, attn, downsample in self.downs:
x = block1(x, t)
x = block2(x, t)
x = attn(x)
h.append(x)
x = downsample(x)
x = self.mid_block1(x, t)
x = self.mid_attn(x)
x = self.mid_block2(x, t)
len_h = len(h) - 1
for block1, block2, attn, upsample in self.ups:
x = ops.concat((x, h[len_h]), 1)
len_h -= 1
x = block1(x, t)
x = block2(x, t)
x = attn(x)
x = upsample(x)
return self.final_conv(x)
```
```python
def linear_beta_schedule(timesteps):
beta_start = 0.0001
beta_end = 0.02
return np.linspace(beta_start, beta_end, timesteps).astype(np.float32)
```
首先，让我们使用 $T=200$ 时间步长的线性计划，并定义我们需要的 $\\β_t$ 中的各种变量，例如方差 $\bar{\alpha}_t$ 的累积乘积。下面的每个变量都只是一维张量，存储从 $t$ 到 $T$ 的值。重要的是，我们还定义了`extract`函数，它将允许我们提取一批适当的 $t$ 索引。
```python
# 扩散200步
timesteps = 200
# 定义 beta schedule
betas = linear_beta_schedule(timesteps=timesteps)
# 定义 alphas
alphas = 1. - betas
alphas_cumprod = np.cumprod(alphas, axis=0)
alphas_cumprod_prev = np.pad(alphas_cumprod[:-1], (1, 0), constant_values=1)
sqrt_recip_alphas = Tensor(np.sqrt(1. / alphas))
sqrt_alphas_cumprod = Tensor(np.sqrt(alphas_cumprod))
sqrt_one_minus_alphas_cumprod = Tensor(np.sqrt(1. - alphas_cumprod))
# 计算 q(x_{t-1} | x_t, x_0)
posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)
p2_loss_weight = (1 + alphas_cumprod / (1 - alphas_cumprod)) ** -0.
p2_loss_weight = Tensor(p2_loss_weight)
def extract(a, t, x_shape):
b = t.shape[0]
out = Tensor(a).gather(t, -1)
return out.reshape(b, *((1,) * (len(x_shape) - 1)))
```
我们将用猫图像说明如何在扩散过程的每个时间步骤中添加噪音。
```python
# 下载猫猫图像
url = 'https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/datasets/image_cat.zip'
path = download(url, './', kind="zip", replace=True)
```
Downloading data from https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/datasets/image_cat.zip (170 kB)
file_sizes: 100%|████████████████████████████| 174k/174k [00:00<00:00, 55.4MB/s]
Extracting zip file...
Successfully downloaded / unzipped to ./
```python
from PIL import Image
image = Image.open('./image_cat/jpg/000000039769.jpg')
base_width = 160
image = image.resize((base_width, int(float(image.size[1]) * float(base_width / float(image.size[0])))))
image.show()
```
噪声被添加到mindspore张量中，而不是Pillow图像。我们将首先定义图像转换，允许我们从PIL图像转换到mindspore张量（我们可以在其上添加噪声），反之亦然。
这些转换相当简单：我们首先通过除以$255$来标准化图像（使它们在 $[0,1]$ 范围内），然后确保它们在 $[-1, 1]$ 范围内。DPPM论文中有介绍到：
> 假设图像数据由 $\{0, 1, ... , 255\}$ 中的整数组成，线性缩放为 $[−1, 1]$ ， 这确保了神经网络反向过程在从标准正常先验 $p(\mathbf{x}_T )$开始的一致缩放输入上运行。
```python
from mindspore.dataset import ImageFolderDataset
image_size = 128
transforms = [
Resize(image_size, Inter.BILINEAR),
CenterCrop(image_size),
ToTensor(),
lambda t: (t * 2) - 1
]
path = './image_cat'
dataset = ImageFolderDataset(dataset_dir=path, num_parallel_workers=cpu_count(),
extensions=['.jpg', '.jpeg', '.png', '.tiff'],
num_shards=1, shard_id=0, shuffle=False, decode=True)
dataset = dataset.project('image')
transforms.insert(1, RandomHorizontalFlip())
dataset_1 = dataset.map(transforms, 'image')
dataset_2 = dataset_1.batch(1, drop_remainder=True)
x_start = next(dataset_2.create_tuple_iterator())[0]
print(x_start.shape)
```
(1, 3, 128, 128)
我们还定义了反向变换，它接收一个包含 $[-1, 1]$ 中的张量，并将它们转回 PIL 图像：
```python
import numpy as np
reverse_transform = [
lambda t: (t + 1) / 2,
lambda t: ops.permute(t, (1, 2, 0)), # CHW to HWC
lambda t: t * 255.,
lambda t: t.asnumpy().astype(np.uint8),
ToPIL()
]
def compose(transform, x):
for d in transform:
x = d(x)
return x
```
让我们验证一下：
```python
reverse_image = compose(reverse_transform, x_start[0])
reverse_image.show()
```
我们现在可以定义前向扩散过程，如本文所示：
```python
def q_sample(x_start, t, noise=None):
if noise is None:
noise = randn_like(x_start)
return (extract(sqrt_alphas_cumprod, t, x_start.shape) * x_start +
extract(sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise) ####对应x0到xt的公式
```
让我们在特定的时间步长上测试它：
```python
def get_noisy_image(x_start, t):
# 添加噪音
x_noisy = q_sample(x_start, t=t)
# 转换为 PIL 图像
noisy_image = compose(reverse_transform, x_noisy[0])
return noisy_image
```
```python
# 设置 time step
t = Tensor([40])
noisy_image = get_noisy_image(x_start, t)
print(noisy_image)
noisy_image.show()
```
<PIL.Image.Image image mode=RGB size=128x128 at 0xFFFED9DF9BE0>
让我们为不同的时间步骤可视化此情况：
```python
import matplotlib.pyplot as plt
def plot(imgs, with_orig=False, row_title=None, **imshow_kwargs):
if not isinstance(imgs[0], list):
imgs = [imgs]
num_rows = len(imgs)
num_cols = len(imgs[0]) + with_orig
_, axs = plt.subplots(figsize=(200, 200), nrows=num_rows, ncols=num_cols, squeeze=False)
for row_idx, row in enumerate(imgs):
row = [image] + row if with_orig else row
for col_idx, img in enumerate(row):
ax = axs[row_idx, col_idx]
ax.imshow(np.asarray(img), **imshow_kwargs)
ax.set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])
if with_orig:
axs[0, 0].set(title='Original image')
axs[0, 0].title.set_size(8)
if row_title is not None:
for row_idx in range(num_rows):
axs[row_idx, 0].set(ylabel=row_title[row_idx])
plt.tight_layout()
```
```python
plot([get_noisy_image(x_start, Tensor([t])) for t in [0, 50, 100, 150, 199]])
```
这意味着我们现在可以定义给定模型的损失函数，如下所示：
```python
# # 计算 q(x_{t-1} | x_t, x_0)
# posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)
# p2_loss_weight = (1 + alphas_cumprod / (1 - alphas_cumprod)) ** -0.
# p2_loss_weight = Tensor(p2_loss_weight)
def p_losses(unet_model, x_start, t, noise=None):
if noise is None:
noise = randn_like(x_start)
x_noisy = q_sample(x_start=x_start, t=t, noise=noise) ####在时间步t破坏图像
predicted_noise = unet_model(x_noisy, t) ###用unet 模型预测 P(Xt-1/Xt)的 噪声分布的均值u
loss = nn.SmoothL1Loss()(noise, predicted_noise)# 加入的真实噪声和预测出来的噪声 做均方差的损失函数 来得到这个步骤中噪声的 距离
loss = loss.reshape(loss.shape[0], -1)
loss = loss * extract(p2_loss_weight, t, loss.shape) ####这里看不懂 请解析
return loss.mean() ####这里看不懂 请解析
```
该代码段是关于如何使用变分自编码器（VAE）或更具体地说，是关于扩散模型（如DDPM，Denoising Diffusion Probabilistic Models）的一部分，用于训练一个模型（这里是unet_model）去学习数据分布并进行去噪。下面是对代码中提到的几个关键点的解释：
1. p_losses函数的目的
此函数旨在计算给定模型（unet_model）在预测噪声分布上的损失，这是在扩散模型训练过程中非常重要的一部分。模型试图学习如何从加噪的数据点x_noisy预测原本的噪声分布，即P(X_{t-1}|X_t)。通过最小化这个损失，模型逐渐学会逆向扩散过程，即从噪声中恢复原始数据。
2. loss = nn.SmoothL1Loss()(noise, predicted_noise)
目的：计算真实噪声（即加到原始数据上的噪声）和模型预测的噪声之间的差距。nn.SmoothL1Loss()是平滑的L1损失函数，它在误差较小的时候接近均方误差（MSE），在误差较大时更像L1损失，这样可以减少大误差对损失函数的敏感度，有助于稳定训练。
操作：真实噪声noise和模型预测的噪声predicted_noise之间的差异被计算出来，作为模型性能的一个度量。
3. loss = loss.reshape(loss.shape[0], -1)
目的：调整形状以便与p2_loss_weight相乘。这一步通常是为了方便后续的操作，确保维度匹配，使得损失可以在每个样本的维度上进行加权。
4. loss = loss * extract(p2_loss_weight, t, loss.shape)
extract函数：虽然代码中没有直接给出extract函数的定义，但从上下文推测，它应该是用于根据当前的时间步t从p2_loss_weight中提取或计算出一个权重，然后将这个权重应用于损失值上。p2_loss_weight是根据时间步调整的权重因子，用来控制不同时间步的损失贡献度。
目的：引入时间依赖的权重因子，这在扩散模型中很常见，用于强调早期时间步（当噪声较高时）的训练，或者是按照特定策略调整训练过程中的重点，以优化模型的学习效率和最终性能。
5. return loss.mean()
目的：计算所有样本损失的平均值。这样做可以得到一个整体的损失指标，用于反向传播和更新模型参数。通过平均损失，可以确保训练过程不会受到单个样本极端损失值的影响，使模型训练更加稳定和均衡。
综上所述，整个函数旨在通过比较模型预测的噪声与实际添加的噪声，计算出损失值，并根据时间步调整这个损失值的权重，最后求所有样本损失的平均值作为训练的反馈信号，指导模型学习如何从噪声中还原出清晰的数据表示。
这段代码是Diffusion Models（扩散模型）中的一部分，用于训练一个Unet模型来预测和还原噪声，进而生成高质量的图像。下面是对您提出疑问部分的详细解析：
### 关于`extract(p2_loss_weight, t, loss.shape)`的解释
首先，这个操作看起来像是自定义的（因为"extract"函数未在提供的代码片段中定义），但可以根据上下文推测其意图。在Diffusion Models的训练过程中，有时会采用时间相关的权重来调整不同时间步长上的损失函数贡献，以优化模型训练，尤其是在涉及噪声预测的设置中。
假设`p2_loss_weight`是一个预计算的时间步依赖的权重张量，形状为`(T,)`，其中`T`是总的 diffusion 步数。这个权重通常用来强调早期或晚期的时间步，以引导模型更好地学习数据分布的不同特性。而`t`是当前时间步的索引，指示了在哪一个时间步上进行的采样和预测。
因此，`extract(p2_loss_weight, t, loss.shape)`这一操作很可能意在根据当前时间步`t`从`p2_loss_weight`中提取相应的权重值，并将其应用于损失函数。具体来说，它可能将一个标量或一维向量（对应于特定时间步的权重）调整成与`loss`形状相匹配的形状，以便逐元素地乘以`loss`。这样做可以确保每个批次中的样本在不同时间步上的损失按照设计的方式被加权。
### 关于`return loss.mean()`的解释
在计算完所有样本在当前时间步的损失后，通过调用`loss.mean()`，目的是计算这些损失值的平均值。这样做有以下几个原因：
1. **简化优化过程**：通过取平均，可以将所有样本的损失合并为一个单一的标量值，这对于梯度下降等优化算法是非常必要的，因为它们通常需要一个标量损失来更新模型参数。
2. **均衡不同样本间的影响**：平均化损失可以确保每个样本对最终损失的贡献是相对平等的，无论批大小如何变化，都能提供一致的训练信号。
3. **稳定训练**：在机器学习中，使用平均损失而不是总损失有助于减少因单个样本异常而导致的训练波动，使得训练过程更加平滑和稳定。
综上所述，`extract(p2_loss_weight, t, loss.shape)`操作是为了根据时间步动态调整损失的权重，以指导模型学习；而`return loss.mean()`则是为了整合所有样本的损失信息，为模型提供一个统一的优化目标。
在Diffusion Models的上下文中，虽然直接的数学公式没有在原始代码段中明确展示，但我们可以基于您的描述和常见实践来推导相关部分的数学表达式。特别是对于`p_losses`函数中的关键步骤，我们可以这样理解：
### 1. 添加噪声的过程（q_sample 函数虽然没给出，但可以依据扩散模型原理描述）
在Diffusion Models中，我们从一个干净的图像 \( x_0 \) 开始，逐步添加噪声生成一个序列 \( x_1, x_2, ..., x_T \)，其中 \( T \) 是总的时间步数。这一过程可以用以下公式概括（尽管实际实现可能涉及更复杂的变换）：
\[ q(x_t|x_0) = \mathcal{N}(x_t; \sqrt{\bar{\alpha}_t}x_0, (1-\bar{\alpha}_t)\mathbf{I}) \]
其中，\( \bar{\alpha}_t = \prod_{s=1}^{t}\alpha_s \) 是从初始时间步到时间步 \( t \) 的噪声衰减系数的乘积，\( \alpha_s \) 对应于每一步的保留信号比例，\( \mathbf{I} \) 是单位矩阵，表明噪声是各向同性的高斯噪声。
### 2. 损失函数计算
- **预测噪声**: Unet模型旨在学习预测给定 \( x_t \) 和时间步 \( t \) 下的噪声分布 \( \epsilon_\theta(x_t, t) \) 的均值，即 \( \hat{\epsilon} = \epsilon_\theta(x_t, t) \)。
- **实际使用的损失函数**: 使用 SmoothL1Loss（也称为Huber Loss），它结合了L1和L2损失的优点，公式如下：
\[ L_{\text{SmoothL1}}(y, \hat{y}) =
\begin{cases}
0.5(y - \hat{y})^2 / \beta & \text{for } |y - \hat{y}| < \beta \\
|y - \hat{y}| - 0.5\beta & \text{otherwise}
\end{cases}
\]
在这里，\( y \) 是真实噪声，\( \hat{y} \) 是预测噪声，而 \( \beta \) 是一个阈值，通常默认设置。
- **时间相关的损失加权** (`p2_loss_weight`): 这个权重项可能是为了根据时间步调整损失的重要性，其公式可能类似于：
\[ w_t = \left(1 + \frac{\bar{\alpha}_t}{1 - \bar{\alpha}_t}\right)^{-\gamma} \]
其中 \( \gamma \) 是一个超参数，用来控制加权的程度。这个表达式可能与你提到的公式不完全一致，但体现了随时间步动态调整权重的思想。
最终，损失函数可以形式化为：
\[ \mathcal{L} = \mathbb{E}_{t \sim [1, T], x_0 \sim p_{data}(x_0), \epsilon \sim \mathcal{N}(0, \mathbf{I})}\left[\left(w_t \cdot | \epsilon - \hat{\epsilon}_\theta(x_t, t)|\right)\right] \]
这里，\( w_t \) 被用于加权每个时间步的损失，以优化模型在不同扩散阶段的表现。
### 3. 平均化损失
\[ \mathcal{L}_{\text{mean}} = \frac{1}{N}\sum_{i=1}^{N}\mathcal{L}_i \]
其中，\( N \) 是批次中的样本数量，\( \mathcal{L}_i \) 是第 \( i \) 个样本的损失，通过平均化得到整个批次的平均损失，作为模型更新的基础。
请注意，具体的实现细节可能会根据实际模型架构和研究者的选择有所不同。
\( \hat{\epsilon} = \epsilon_\theta(x_t, t) \) 是对应于\( \hat{\epsilon} (x_{t-1} | x_t)的吗 也就是unet 网络实际上预测的是Xt步到Xt-1步骤的噪声 对吗 是不是公式有体现？
实际上，公式 \( \hat{\epsilon} = \epsilon_\theta(x_t, t) \) 表示UNet网络预测的是在时间步 \( t \) 给定的噪声 \( \epsilon \)（这个噪声是理论上添加到 \( x_0 \) 上以得到 \( x_t \) 的噪声）。也就是说，网络尝试直接预测给定当前带噪图像 \( x_t \) 和时间步指示器 \( t \) 条件下的噪声分量 \( \epsilon \)。
至于 \( \hat{\epsilon} (x_{t-1} | x_t) \)，这个表述有些误导性。在标准的扩散模型框架里，我们并不直接预测从 \( x_t \) 到 \( x_{t-1} \) 的噪声，而是预测在时间步 \( t \) 的条件下，原始数据 \( x_0 \) 被添加的噪声 \( \epsilon \)。这是因为 diffusion process 是一个正向过程（从纯净图像逐步添加噪声直到完全随机），而训练UNet模型是为了在反向过程（从噪声中恢复图像）中表现良好，它通过学习如何估计在每个时间步上添加的噪声来间接实现这一点。
因此，正确的理解是，UNet模型的输出 \( \hat{\epsilon} \) 应该被解释为对"在时间点 \( t \)，如果从 \( x_0 \) 加入噪声得到 \( x_t \)，那么这部分噪声是多少"的预测，而不是直接预测从 \( x_t \) 如何变化到 \( x_{t-1} \)。后者实际上是通过迭代过程和已知的噪声分布逐渐逆向去噪实现的，而非网络直接预测两步之间的噪声差。
我们并不直接预测从 \( x_t \) 到 \( x_{t-1} \) 的噪声，而是预测在时间步 \( t \) 的条件下，原始数据 \( x_0 \) 被添加的噪声 \( \epsilon \)。 公式怎么体现这个过程
理解这一点的关键在于如何将预测噪声的过程与从 \( x_t \) 到 \( x_{t-1} \) 的逆向扩散联系起来，尽管直接的数学公式在原始提问中没有完全展现，我们可以构建起这两者之间的逻辑桥梁。
### 预测噪声的公式
直接反映UNet模型预测噪声任务的公式是 \( \hat{\epsilon} = \epsilon_\theta(x_t, t) \)，其中 \( \hat{\epsilon} \) 是模型预测的噪声，\( x_t \) 是当前噪声水平下的图像，\( t \) 是时间步，而 \( \theta \) 表示模型参数。这个公式表明，给定当前时间步的图像 \( x_t \) 和时间步指示 \( t \)，模型试图预测最初添加到纯净图像 \( x_0 \) 上的噪声分量 \( \epsilon \)。
### 逆向过程的数学联系
虽然直接预测 \( x_{t-1} \) 的噪声变化没有显式公式化，但其背后的原理与以下步骤紧密相关：
1. **噪声预测**：如前所述，模型预测 \( \hat{\epsilon} \)。
2. **逆向一步**：为了从 \( x_t \) 回复到 \( x_{t-1} \)，理论上我们利用以下公式（或其变形，具体取决于模型的设定）：
\[ x_{t-1} = \frac{1}{\sqrt{\alpha_t}}(x_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}}\hat{\epsilon}) \]
这里，\( \alpha_t = 1 - \beta_t \) 是保留信号的比例，\( \bar{\alpha}_t = \prod_{s=1}^{t}\alpha_s \) 是累积至今的信号保留比例。这个公式实质上是根据当前时刻的图像 \( x_t \) 和预测的噪声 \( \hat{\epsilon} \)，以及已知的扩散参数，逆推一步到前一时间步的图像表示。
### 整体流程体现
整个流程虽然没有直接通过一个公式完整展现从 \( x_t \) 到 \( x_{t-1} \) 的噪声预测，但通过预测 \( x_0 \) 在时间步 \( t \) 所承受的噪声 \( \epsilon \)，并利用该预测结合扩散过程的逆向方程，模型实际上是在逐步学习如何去除噪声，从而间接地实现从当前噪声图像 \( x_t \) 回复到更早时间步 \( x_{t-1} \) 或最终到 \( x_0 \) 的过程。因此，虽然预测噪声本身不直接指明 \( x_t \) 到 \( x_{t-1} \) 的噪声变化，但它是实现这一逆向过程中的核心步骤之一。
公式 \( \hat{\epsilon} = \epsilon_\theta(x_t, t) \) 表示UNet网络预测的是在时间步 \( t \) 给定的噪声 \( \epsilon \)（这个噪声是理论上添加到 \( x_0 \) 上以得到 \( x_t \) 的噪声） 公式具体是怎么推导出来的
公式 \( \hat{\epsilon} = \epsilon_\theta(x_t, t) \) 实际上是基于Diffusion Model的训练目标和设计逻辑，而非直接从数学推导而来。不过，我们可以理解这个公式的含义及其背后的思想：
### Diffusion Model的基本思想
在Diffusion Models中，目标是从一个容易生成的噪声分布中采样，然后通过一系列逆扩散步骤生成高质量的数据样本（如图像）。这个过程可以看作是两个相反的过程：正向扩散（Forward Process）和逆向扩散（Reverse Process）。
#### 正向扩散（Forward Process）
正向扩散过程通过逐渐添加噪声到原始数据 \( x_0 \) 上，经过 \( T \) 个时间步后，得到完全噪声化的数据 \( x_T \)。每一时间步 \( t \)，噪声添加的概率模型通常被建模为高斯过程，形式化为：
\[ q(x_t|x_{t-1}) = \mathcal{N}(x_t; \sqrt{1-\beta_t}x_{t-1}, \beta_t\mathbf{I}) \]
其中，\( \beta_t \) 是控制每一时间步噪声量的超参数，\( \mathbf{I} \) 是单位矩阵。
从 \( x_0 \) 到 \( x_t \) 的过程可以递归地表示，最终可以推导出 \( x_t \) 相对于 \( x_0 \) 的表达式，其中包含了一个累计噪声项，这个噪声项实质上就是 \( \epsilon \)。
#### 逆向扩散（Reverse Process）
逆向扩散的目标是从 \( x_T \) 开始，逐步减少噪声，直至恢复出 \( x_0 \)。由于直接计算逆过程的概率分布非常复杂，Diffusion Models采用变分推理方法，通过训练神经网络（如UNet）来近似逆过程的每一步。
### UNet预测噪声的逻辑
UNet模型 \( \epsilon_\theta \) 被设计来预测在时间步 \( t \) 的噪声分量 \( \epsilon \)，这个预测是基于当前观测到的带噪图像 \( x_t \) 和时间步信息 \( t \)。具体来说，模型试图学习：
\[ \hat{\epsilon} = \epsilon_\theta(x_t, t) \approx \epsilon \]
这里的 \( \hat{\epsilon} \) 是模型预测的噪声，而 \( \epsilon \) 是理论上添加到 \( x_0 \) 上以生成当前 \( x_t \) 的噪声。
这个公式的提出并非直接数学推导的结果，而是基于Diffusion Models设计的逻辑推断：要从噪声中逆向生成数据，我们需要知道每一步如何添加或移除噪声，而UNet通过学习预测每一步的噪声分布来逼近这一过程。这一过程的训练依据是最大化似然，通过比较模型预测的噪声与实际添加的噪声（或通过其他损失函数，如MSE、SmoothL1Loss等），逐步优化模型参数 \( \theta \) 以达到更好的逆向扩散效果。
```python
```
```python
```
```python
```
`denoise_model`将是我们上面定义的U-Net。我们将在真实噪声和预测噪声之间使用Huber损失。
## 数据准备与处理
在这里我们定义一个正则数据集。数据集可以来自简单的真实数据集的图像组成，如Fashion-MNIST、CIFAR-10或ImageNet，其中线性缩放为 $[−1, 1]$ 。
每个图像的大小都会调整为相同的大小。有趣的是，图像也是随机水平翻转的。根据论文内容：我们在CIFAR10的训练中使用了随机水平翻转；我们尝试了有翻转和没有翻转的训练，并发现翻转可以稍微提高样本质量。
本实验我们选用Fashion_MNIST数据集，我们使用download下载并解压Fashion_MNIST数据集到指定路径。此数据集由已经具有相同分辨率的图像组成，即28x28。
```python
# 下载MNIST数据集
url = 'https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/datasets/dataset.zip'
path = download(url, './', kind="zip", replace=True)
```
Downloading data from https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/datasets/dataset.zip (29.4 MB)
file_sizes: 100%|██████████████████████████| 30.9M/30.9M [00:00<00:00, 65.3MB/s]
Extracting zip file...
Successfully downloaded / unzipped to ./
```python
from mindspore.dataset import FashionMnistDataset
image_size = 28
channels = 1
batch_size = 16
fashion_mnist_dataset_dir = "./dataset"
dataset = FashionMnistDataset(dataset_dir=fashion_mnist_dataset_dir, usage="train", num_parallel_workers=cpu_count(), shuffle=True, num_shards=1, shard_id=0)
```
接下来，我们定义一个transform操作，将在整个数据集上动态应用该操作。该操作应用一些基本的图像预处理：随机水平翻转、重新调整，最后使它们的值在 $[-1,1]$ 范围内。
```python
transforms = [
RandomHorizontalFlip(),
ToTensor(),
lambda t: (t * 2) - 1
]
dataset = dataset.project('image')
dataset = dataset.shuffle(64)
dataset = dataset.map(transforms, 'image')
dataset = dataset.batch(16, drop_remainder=True)
```
```python
x = next(dataset.create_dict_iterator())
print(x.keys())
```
dict_keys(['image'])
### 采样
由于我们将在训练期间从模型中采样（以便跟踪进度），我们定义了下面的代码。采样在本文中总结为算法2：
从扩散模型生成新图像是通过反转扩散过程来实现的：我们从$T$开始，我们从高斯分布中采样纯噪声，然后使用我们的神经网络逐渐去噪（使用它所学习的条件概率），直到我们最终在时间步$t = 0$结束。如上图所示，我们可以通过使用我们的噪声预测器插入平均值的重新参数化，导出一个降噪程度较低的图像
$\mathbf{x}_{t-1 }$。请注意，方差是提前知道的。
理想情况下，我们最终会得到一个看起来像是来自真实数据分布的图像。
下面的代码实现了这一点。
```python
def p_sample(model, x, t, t_index):
betas_t = extract(betas, t, x.shape)
sqrt_one_minus_alphas_cumprod_t = extract(
sqrt_one_minus_alphas_cumprod, t, x.shape
)
sqrt_recip_alphas_t = extract(sqrt_recip_alphas, t, x.shape)
model_mean = sqrt_recip_alphas_t * (x - betas_t * model(x, t) / sqrt_one_minus_alphas_cumprod_t)
if t_index == 0:
return model_mean
posterior_variance_t = extract(posterior_variance, t, x.shape)
noise = randn_like(x)
return model_mean + ops.sqrt(posterior_variance_t) * noise
def p_sample_loop(model, shape):
b = shape[0]
# 从纯噪声开始
img = randn(shape, dtype=None)
imgs = []
for i in tqdm(reversed(range(0, timesteps)), desc='sampling loop time step', total=timesteps):
img = p_sample(model, img, ms.numpy.full((b,), i, dtype=mstype.int32), i)
imgs.append(img.asnumpy())
return imgs
def sample(model, image_size, batch_size=16, channels=3):
return p_sample_loop(model, shape=(batch_size, channels, image_size, image_size))
```
请注意，上面的代码是原始实现的简化版本。
# 详细的代码
# # ------------------------------------------ sampling ------------------------------------------#
# def predict_start_from_noise(self, x_t, t, noise):
# '''
# if self.predict_epsilon, model output is (scaled) noise;
# otherwise, model predicts x0 directly
# '''
# if self.predict_epsilon:
# return (
# extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
# extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise
# )
# else:
# return noise
# def q_posterior(self, x_start, x_t, t):
# posterior_mean = (
# extract(self.posterior_mean_coef1, t, x_t.shape) * x_start +
# extract(self.posterior_mean_coef2, t, x_t.shape) * x_t
# )
# posterior_variance = extract(self.posterior_variance, t, x_t.shape)
# posterior_log_variance_clipped = extract(self.posterior_log_variance_clipped, t, x_t.shape)
# return posterior_mean, posterior_variance, posterior_log_variance_clipped
# def p_mean_variance(self, x, t, s):
# x_recon = self.predict_start_from_noise(x, t=t, noise=self.model(x, t, s))
# if self.clip_denoised:
# x_recon.clamp_(-self.max_action, self.max_action)
# else:
# assert RuntimeError()
# model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
# return model_mean, posterior_variance, posterior_log_variance
# # @torch.no_grad()
# def p_sample(self, x, t, s):
# b, *_, device = *x.shape, x.device
# model_mean, _, model_log_variance = self.p_mean_variance(x=x, t=t, s=s)
# noise = torch.randn_like(x)
# # no noise when t == 0
# nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
# return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
# # @torch.no_grad()
# def p_sample_loop(self, state, shape, verbose=False, return_diffusion=False):
# device = self.betas.device
# batch_size = shape[0]
# x = torch.randn(shape, device=device)
# if return_diffusion: diffusion = [x]
# progress = Progress(self.n_timesteps) if verbose else Silent()
# for i in reversed(range(0, self.n_timesteps)):
# timesteps = torch.full((batch_size,), i, device=device, dtype=torch.long)
# x = self.p_sample(x, timesteps, state)
# progress.update({'t': i})
# if return_diffusion: diffusion.append(x)
# progress.close()
# if return_diffusion:
# return x, torch.stack(diffusion, dim=1)
# else:
# return x
# # @torch.no_grad()
# def sample(self, state, *args, **kwargs):
# batch_size = state.shape[0]
# shape = (batch_size, self.action_dim)
# action = self.p_sample_loop(state, shape, *args, **kwargs)
# return action.clamp_(-self.max_action, self.max_action)
# # ------------------------------------------ training ------------------------------------------#
# def q_sample(self, x_start, t, noise=None):
# if noise is None:
# noise = torch.randn_like(x_start)
# sample = (
# extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
# extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
# )
# return sample
# def p_losses(self, x_start, state, t, weights=1.0):
# noise = torch.randn_like(x_start)
# x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
# x_recon = self.model(x_noisy, t, state)
# assert noise.shape == x_recon.shape
# if self.predict_epsilon:
# loss = self.loss_fn(x_recon, noise, weights)
# else:
# loss = self.loss_fn(x_recon, x_start, weights)
# return loss
# def loss(self, x, state, weights=1.0):
# batch_size = len(x)
# t = torch.randint(0, self.n_timesteps, (batch_size,), device=x.device).long()
# return self.p_losses(x, state, t, weights)
# def forward(self, state, *args, **kwargs):
# return self.sample(state, *args, **kwargs)
```python
```
```python
```
```python
```
```python
```
```python
```
## 训练过程
下面，我们开始训练吧！
```python
# 定义动态学习率
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_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
# 定义优化器
optimizer = nn.Adam(unet_model.trainable_params(), learning_rate=lr)
loss_scaler = DynamicLossScaler(65536, 2, 1000)
# 定义前向过程
def forward_fn(data, t, noise=None):
loss = p_losses(unet_model, data, t, noise)
return loss
# 计算梯度
grad_fn = ms.value_and_grad(forward_fn, None, optimizer.parameters, has_aux=False)
# 梯度更新
def train_step(data, t, noise):
loss, grads = grad_fn(data, t, noise)
optimizer(grads)
return loss
```
```python
randint(0, timesteps, (batch_size,), dtype=ms.int32)
```
Tensor(shape=[16], dtype=Int32, value= [ 52, 77, 47, 5, 56, 165, 181, 176, 116, 26, 182, 131, 14, 140, 29, 32])
```python
import time
# 由于时间原因，epochs设置为1，可根据需求进行调整
epochs = 1
for epoch in range(epochs):
begin_time = time.time()
for step, batch in enumerate(dataset.create_tuple_iterator()):
unet_model.set_train()
batch_size = batch[0].shape[0]
t = randint(0, timesteps, (batch_size,), dtype=ms.int32)
noise = randn_like(batch[0])
loss = train_step(batch[0], t, noise) ###t 是一个shape=[16] 包括了16个时间步骤的噪声预测
if step % 500 == 0:
print(" epoch: ", epoch, " step: ", step, " Loss: ", loss)
end_time = time.time()
times = end_time - begin_time
print("training time:", times, "s")
print("sample-------------------------------------------")
# 展示随机采样效果
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!")
```
```python
print("yangge DDPM05 训练过程和采样过程 2024 06 28")
```
## 推理过程（从模型中采样）
要从模型中采样，我们可以只使用上面定义的采样函数：
```python
# 采样64个图片
unet_model.set_train(False)
samples = sample(unet_model, image_size=image_size, batch_size=64, channels=channels)
```
```python
# 展示一个随机效果
random_index = 5
plt.imshow(samples[-1][random_index].reshape(image_size, image_size, channels), cmap="gray")
```
可以看到这个模型能产生一件衣服！
请注意，我们训练的数据集分辨率相当低（28x28）。
我们还可以创建去噪过程的gif：
```python
import matplotlib.animation as animation
random_index = 53
fig = plt.figure()
ims = []
for i in range(timesteps):
im = plt.imshow(samples[i][random_index].reshape(image_size, image_size, channels), cmap="gray", animated=True)
ims.append([im])
animate = animation.ArtistAnimation(fig, ims, interval=50, blit=True, repeat_delay=100)
animate.save('diffusion.gif')
plt.show()
```
在扩散模型的采样阶段，也就是生成新的数据点（如图像）的过程中，UNet模型的参数被用来逆向执行扩散过程。这个阶段通常涉及一系列迭代的去噪步骤，从高度噪声化的数据开始，逐步恢复到清晰的原始数据。具体来说，采样阶段的操作可以分为以下几个关键步骤：
1. **初始化**: 采样过程通常从生成一个完全随机的噪声图像\( x_T \)开始，这个图像没有携带任何有用的信息。
2. **逐步去噪**: 接下来，按照时间步从\( T \)递减到0的顺序，应用UNet模型来逐步减少噪声。对于每一个时间步\( t \)，模型预测噪声\( \hat{\epsilon}(x_t, t) \)，并基于这个预测和当前时间步的图像\( x_t \)计算出更清晰的图像\( x_{t-1} \)。这个计算涉及到一个由模型参数和时间步决定的去噪公式。
- 去噪公式可能类似于：
\[
x_{t-1} = \frac{1}{\sqrt{\alpha_t}} \left( x_t - \frac{(1 - \alpha_t)}{\sqrt{1 - \bar{\alpha}_t}} \hat{\epsilon}(x_t, t) \right)
\]
其中，\( \alpha_t \)和\( \bar{\alpha}_t \)是预定义的衰减系数，用于控制噪声的减少程度。
3. **迭代**: 上述去噪步骤会重复执行，直到\( t \)达到0。在每次迭代中，模型都会基于当前的\( x_t \)和时间步\( t \)预测噪声\( \hat{\epsilon}(x_t, t) \)，并使用该预测更新图像到下一个更清晰的状态\( x_{t-1} \)。
4. **终止**: 当\( t = 0 \)时，最终的图像\( x_0 \)被认为是从随机噪声开始通过反向扩散过程生成的新样本。
### 重要概念
- **时间步**: 每个时间步\( t \)不仅影响噪声的预测，还决定了去噪的强度和方式。随着时间步的减少，图像逐渐从噪声中恢复细节。
- **迭代过程**: 采样不是一步到位的，而是通过一系列小步骤逐步逼近清晰图像。每个步骤都是基于模型对噪声的预测和当前时间步的去噪策略。
- **预训练模型**: 在采样阶段使用的UNet模型参数已经在训练阶段通过最大化后验概率的方式优化过，即学习了如何在每个时间步有效地预测和去除噪声
在变分自编码器（VAE）的上下文中，均值 \(\mu\) 和方差（或者更常见的是对数方差 \(\log\sigma^2\)）并不是单个标量值，而是向量。这是因为潜在空间（latent space）是多维的，每个维度都有其自己的均值和方差。
例如，如果我们说潜在空间是64维的，那么 \(\mu\) 和 \(\log\sigma^2\) 都将是64维的向量。每个维度的 \(\mu\) 和 \(\log\sigma^2\) 的值表示在潜在空间的该特定维度上的分布参数。因此，\(\mu\) 和 \(\log\sigma^2\) 分别表示潜在空间中各维度的均值和对数方差。
当我们说采样 \(z\) 时，我们从一个多维的高斯分布中采样，其中每个维度的高斯分布具有不同的均值和方差，这些参数由 \(\mu\) 和 \(\log\sigma^2\) 向量给出。因此，\(z\) 也是一个64维的向量，其中每个元素都是从潜在空间中对应维度的高斯分布中独立采样的。
这种设计允许模型学习输入数据的复杂结构，其中不同的潜在因素可能独立变化。通过给每个潜在维度分配独立的均值和方差，VAE 能够捕获数据集中的不同特征和模式。这就是为什么在高维潜在空间中，均值和方差都是向量的原因。
1248
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
