DiffusionModel-Classifier Guidance《Diffusion Models Beat GANs on Image Synthesis》原理＋核心代码解读

论文
https://arxiv.org/pdf/2105.05233
代码
https://github.com/openai/guided-diffusion
参考了一下大佬的博客：
classifier guided diffusion

GAN牺牲多样性从而提高了介绍，
文章提到了一种名为“Classifier Guidance”的方法，用于在扩散模型的采样过程中引入分类器的梯度，以指导生成过程。这种方法允许在生成图像时根据分类器的梯度来调整模型的采样过程，从而实现对特定类别的生成。

• 分类器的作用：
  • 分类器 p(y|x) 可以根据输入图像 x 预测它属于类别 y 的概率。
  • 分类器不仅用于分类，还用于指导生成模型生成特定类别的图像
• $P_{\Phi }(y|x_t,t)$:我们可以在噪声图像$x_t$上训练分类器
• ${\nabla }_{x_t}\log P_{\Phi }(y|x_t,t)$:引导扩散采样过程朝向任意类标签 y

先验知识

多元高斯分布对数似然函数

多元高斯分布（Multivariate Gaussian Distribution）是指具有多个随机变量的高斯分布。在多元高斯分布中，随机变量是一个多维向量，而不是单个标量。多元高斯分布可以用来描述多维数据的分布特征，比如多维空间中的数据点分布情况。

多元高斯分布的概率密度函数可以表示为：

$f(x|\mu ,\Sigma )=\frac{1}{(2\pi {)}^{n/2}|\Sigma {|}^{1/2}}\exp \left(-\frac{1}{2}(x-\mu {)}^T{\Sigma }^{-1}(x-\mu )\right)$

其中$x$ 是一个$n$维向量，$\mu$是多元高斯分布的均值向量，$\Sigma$ 是多元高斯分布的协方差矩阵，$n$ 是向量的维度。

$$\begin{array}{rl}l(\mu ,\Sigma |{\mathbf{x}}^{(\mathbf{i})})& =\log \prod _{i=1}^m{f}_{\mathbf{X}(\mathrm{i})}({\mathbf{x}}^{(\mathrm{i})}|\mu ,\mathbf{\Sigma })\\ & =\log \prod _{i=1}^m\frac{1}{(2\pi {)}^{p/2}{\left|\Sigma \right|}^{1/2}}\exp \left(-\frac{1}{2}({\mathbf{x}}^{(\mathrm{i})}-\mu {)}^{\mathrm{T}}{\mathbf{\Sigma }}^{-1}({\mathbf{x}}^{(\mathrm{i})}-\mu )\right)\\ & =\sum _{i=1}^m\left(-\frac{p}{2}\log (2\pi )-\frac{1}{2}\log |\Sigma |-\frac{1}{2}({\mathbf{x}}^{(\mathrm{i})}-\mu {)}^{\mathrm{T}}{\mathbf{\Sigma }}^{-1}({\mathbf{x}}^{(\mathrm{i})}-\mu )\right)\\ l(\mu ,\Sigma ;)& =-\frac{mp}{2}\log (2\pi )-\frac{m}{2}\log |\Sigma |-\frac{1}{2}\sum _{i=1}^m({\mathbf{x}}^{(\mathrm{i})}-\mu {)}^{\mathrm{T}}{\mathbf{\Sigma }}^{-1}({\mathbf{x}}^{(\mathrm{i})}-\mu )\end{array}$$

• $\mu$是多元高斯分布的均值向量
• $\Sigma$是多元高斯分布的协方差矩阵
• X = { x 1 , x 2 , . . . , x N } X = \{x_1, x_2, ..., x_N\} X={x1​,x2​,...,xN​}是观测数据集
• $n$ 是数据的维度。
• 多元高斯分布的协方差矩阵$\Sigma$是对称和半正定矩阵

生成式模型的效果

两大目标：真实性（Quality）和多样性（Diversity）

Inception score

Inception Score基于Inception-V3模型的分类概率来评估生成照片的质量，通过使用预训练的卷积神经网络（通常是Inception网络）来评估生成图像的质量和多样性。
$$\mathrm{KL}\mathrm{divergence}=\mathrm{p}(y|x)\ast (\log (\mathrm{p}(y|x))-\log (\mathrm{p}(\mathrm{y})))$$

• p（y|x）:表示在给定条件 𝑥下，观测到的概率 𝑦
• p(y) 表示概率 𝑦的边际概率
• 对KL散度对所有类别求和再取平均值，并且取一个e指数，即可得到Inception Score。一般生成5000张照片S的值在0~1000范围内。
• 作者希望p(y|x)应该具有低熵(即越真实），p(y)应该具有高熵即越多样，因此，IS值越大越好
• 缺点:缺乏跟真实照片之间的比较;缺乏类内多样性，例如每个类别只产生一张一模一样的照片，IS一样很高

# 用代码实现IS

def calculate_inception_score(p_yx, eps=1E-16):
    # calculate p(y)
    p_y = expand_dims(p_yx.mean(axis=0), 0)#对所有样本求平均
    # kl divergence for each image
    kl_d = p_yx * (log(p_yx + eps) - log(p_y + eps))
    # sum over classes
    sum_kl_d = kl_d.sum(axis=1)#对类别求和
    # average over images，对所有图像的散度取平均值
    avg_kl_d = mean(sum_kl_d)
    # undo the logs
    is_score = exp(avg_kl_d)
    return is_score

Frechlet Inception Distance（FID）

FID是基于Inception-V3模型（预训练好的图像分类模型）的feature vectors来计算真实图片与生成图片之间的距离，用高斯分布来表示，FID就是计算两个分布之间的Wasserstein-2距离。将真实图片和预测图片分别经过Inception模型中，得到2048维度（特征的维度）的embedding vector。把生成和真实的图片同时放入Inception-V3中，然后将feature vectors取出来用于比较。
$$d^2=||{\mu }_1-{\mu }_2|{|}^2+Tr(C_1+C_2-2\ast \sqrt{(C_1\ast C_2)})$$

• μ1，μ2为均值，c1,c2为协方差，Tr为矩阵的迹
• FID越低，说明预测分布越接近于真实的分布
• 可以评估类内多样性，例如每个类别只产生一张一模一样的照片，FID会比较高，也就意味着评估效果比较差

# 用代码实现FID

def calculate_fid(act1, act2):
    # 计算mean和covariance统计量
    #act1，act2是两个隐变量（分别对应真实和预测的照片送入inception-v3得到的隐变量）
    mu1, sigma1 = act1.mean(axis=0), cov(act1, rowvar=False)
    mu2, sigma2 = act2.mean(axis=0), cov(act2, rowvar=False)
    
    # 计算mean之间的平方差之和
    ssdiff = numpy.sum((mu1 - mu2)**2.0)
    
    # 计算cov之积的平方根
    covmean = sqrtm(sigma1.dot(sigma2))
    
    # 检查从sqrt中移除的复数部分
    if iscomplexobj(covmean):
        covmean = covmean.real

    # 计算score
    fid = ssdiff + trace(sigma1 + sigma2 - 2.0 * covmean)
    return fid

Precision和Recall

评估多样性和分布的覆盖度

Precision：预测分布/真实分布
Recall：真实分布/预测分布

可学习的方差${\Sigma }_{\theta }(x_t,t)=\exp (v\log {\beta }_t+(1-v)\log {\tilde{\beta }}_t)$

结构的优化

Unet的改变

• 增加深度的同时保持模型大小相对恒定。
• 增加attention heads的数量。
• 使用在32x32、16x16和8x8分辨率上的attention，而不仅仅是16x16分辨率。
• 使用BigGAN的残差块进行激活的上采样和下采样。
• 对残差连接进行重缩放，使用 $\frac{1}{\sqrt{2}}$。

Adaptive Group Normalization 自适应GN

Adaptive Group Normalization在组标准化操作之后将时间步长和类嵌入合并到每个残差块中
$$AdaGN(h,y)=y_{s}GroupNorm(h)+y_b$$

• h 是残差块中第一层卷积后的中间激活
• y = [ys, yb] 是通过时间步和类别嵌入的线性投影得到的。这种层将时间步和类别信息整合到每个残差块中，有助于模型学习到更丰富的特征表示。

基于分类器的条件采样算法的原理与效果

$\hat{q}$的逆向扩散

保留原来的扩散模型不变，额外引入classifier.
先定义一个马尔科夫类型的扩散过程$\hat{q}$,并$\hat{q}(y|x_0)$都有对应便签，即每个训练样本都有对应标签：
$$\begin{array}{rl}\hat{q}(x_0)& :=q(x_0)\\ \hat{q}(y|x_0)& :=\mathrm{Known}\mathrm{labels}\mathrm{per}\mathrm{sample}\\ \hat{q}(x_{t+1}|x_t,y)& :=q(x_{t+1}|x_t)\\ \hat{q}(x_{1:T}|x_0,y)& :=\prod _{t=1}^T\hat{q}(x_t|x_{t-1},y)\end{array}$$

• 训练样本不变
• 训练样本有标签
• 定义$\hat{q}$与q有一样的加噪过程
• $\hat{q}(x_{1:T}|x_0,y)$马尔科夫类型

证明：分别证明$\hat{q}$加噪条件分布、联合分布和边缘分布，在不加y条件的情况下，$\hat{q}$与$q$,的表现相同;并且进一步表明逆扩散条件分布也相同：

加噪条件分布相同：
$$\begin{array}{rl}\hat{q}(x_{t+1}|x_t)& \stackrel{\text{全概率公式}}{\to }{\int }_y\hat{q}(x_{t+1},y|x_t)dy\\ & \stackrel{\text{联合概率}}{\to }{\int }_y\hat{q}(x_{t+1}|x_t,y)\hat{q}(y|x_t)dy\\ & \stackrel{\text{代入}\hat{q}的定义}{\to }{\int }_{y}q(x_{t+1}|x_t)\hat{q}(y|x_t)dy\\ & =q(x_{t+1}|x_t){\int }_y\hat{q}(y|x_t)dy\\ & =q(x_{t+1}|x_t)\\ & =\hat{q}(x_{t+1}|x_t,y)\end{array}$$
联合概率分布相同：
$$\begin{array}{rl}\hat{q}(x_{1:T}|x_0)& ={\int }_y\hat{q}(x_{1:T},y|x_0)dy\\ & ={\int }_y\hat{q}(y|x_0)\hat{q}(x_{1:T}|x_0,y)dy\\ & ={\int }_y\hat{q}(y|x_0)\prod _{t=1}^T\hat{q}(x_t|x_{t-1},y)dy\\ & ={\int }_y\hat{q}(y|x_0)\prod _{t=1}^{T}q(x_t|x_{t-1})dy\\ & =\prod _{t=1}^{T}q(x_t|x_{t-1}){\int }_y\hat{q}(y|x_0)dy\\ & =\prod _{t=1}^{T}q(x_t|x_{t-1})\\ & =q(x_{1:T}|x_0)\end{array}$$

边缘分布相同：
$$\begin{array}{rl}\hat{q}(x_t)& ={\int }_{x_{0:t-1}}\hat{q}(x_0,...,x_t)d{x}_{0:t-1}\\ & ={\int }_{x_{0:t-1}}\hat{q}(x_0)\hat{q}(x_1,...,x_t|x_0)d{x}_{0:t-1}\\ & ={\int }_{x_{0:t-1}}q(x_0)q(x_1,...,x_t|x_0)d{x}_{0:t-1}\\ & ={\int }_{x_{0:t-1}}q(x_0,...,x_t)d{x}_{0:t-1}\\ & =q(x_t)\end{array}$$
注意：$p_{\theta ,\varphi }(y|x_t,t)$=$p_{\varphi }(y|x_t,t)$and ${ϵ}_{\theta }(x_t,t)={ϵ}_{\theta }(x_t)$ for brevity

逆向过程：可以知道任意时刻t的分布$x_t$和$x_{t-1}$无关
$$\begin{array}{rl}\hat{q}(y|x_t,x_{t+1})& \stackrel{\text{贝叶斯公式},把x_t看成是固定的}{\to }\hat{q}(x_{t+1}|x_t,y)\frac{\hat{q}(y|x_t)}{\hat{q}(x_{t+1}|x_t)}\\ & \stackrel{加噪条件分布相同}{\to }\hat{q}(x_{t+1}|x_t)\frac{\hat{q}(y|x_t)}{\hat{q}(x_{t+1}|x_t)}\\ & =\hat{q}(y|x_t)\end{array}$$

给定条件($x_{t+1}$,y)下的逆扩散条件分布：

$$\begin{array}{rl}\hat{q}(x_t|x_{t+1},y)& \stackrel{把条件分布写成联合分布除以边缘分布}{\to }\frac{\hat{q}(x_t,x_{t+1},y)}{\hat{q}(x_{t+1},y)}\\ & \stackrel{进一步拆分联合分布}{\to }\frac{\hat{q}(x_t,x_{t+1},y)}{\hat{q}(y|x_{t+1})\hat{q}(x_{t+1})}\\ & \stackrel{进一步拆分边缘分布}{\to }\frac{\hat{q}(x_t|x_{t+1})\hat{q}(y|x_t,x_{t+1})\hat{q}(x_{t+1})}{\hat{q}(y|x_{t+1})\hat{q}(x_{t+1})}\\ & \stackrel{消元}{\to }\frac{\hat{q}(x_t|x_{t+1})\hat{q}(y|x_t,x_{t+1})}{\hat{q}(y|x_{t+1})}\\ & \stackrel{与x_{t+1}无关（之前推导的）}{\to }\frac{\hat{q}(x_t|x_{t+1})\hat{q}(y|x_t)}{\hat{q}(y|x_{t+1})}\\ & \stackrel{边缘分布相同}{\to }\boxed{\frac{q(x_t|x_{t+1})\hat{q}(y|x_t)}{\hat{q}(y|x_{t+1})}}\end{array}$$

• q:训练好的扩散过程
• $\hat{q}(y|x_{t+1})$:可以看成一个常量，因为它与$x_t$无关
• $\hat{q}(y|x_{t+1})\sim Zq(x_t|x_{t+1})\hat{q}(y|x_t)$
  • Z:归一化常数
  • 已经有$q(x_t|x_{t+1})$的近似估计，$p_{\theta }(x_t|x_{t+1})$
• $\hat{q}(y|x_t)$:只有这个未知。可以用一个在从$q(x_t)$采样的$x_t$训练的classifier$p_{\varphi }(y|x_t)$

Conditional Reverse Noising Process

在生成过程中加入一个额外的标签 𝑦 来引导生成的样本从而将无条件的扩散模型变成有条件的扩散模型，已经证明：

$$p_{\theta ,\varphi }(x_t|x_{t+1},y)=Z{p}_{\theta }(x_t|x_{t+1})p_{\varphi }(y|x_t)$$

• z：归一化常数，保证所有的可能xt概率之和为1
• y：标签
• $p_{\theta ,\varphi }(x_t|x_{t+1},y)$，给定条件y下的条件概率分布
• φ，θ：分类器参数和模型参数
• $p_{\varphi }(y|x_t)$: 这是在给定数据 $x_t$ 的条件下类别标签 ( y ) 的概率，由分类器决定，参数由 ( \phi ) 参数化。分类器用于引导扩散过程生成与特定类别一致的数据。
• $p_{\theta }(x_t|x_{t+1})$: 这是扩散模型的转移概率，由参数 $\theta$ 参数化。它代表了在没有考虑类别标签的情况下，给定后续时间步的数据 $x_{t+1}$，时间步$t$ 的数据的概率。

之前的diffusion model采样的Gaussian sidtribution：
$$\begin{array}{rl}{p}_{\theta }(x_t|x_{t+1})& =\mathcal{N}(\mu ,\Sigma )\\ \log p_{\theta }(x_t|x_{t+1})& =-\frac{1}{2}(x_t-\mu {)}^T{\Sigma }^{-1}(x_t-\mu )+C\end{array}$$

设$p_{\varphi }$可微，那么就可以进行泰勒展开，展开到第二项，$\log p_{\varphi }(y|x_t)$相对于$\Sigma$的曲率很小，我们的泰勒展开展开到一阶：
$$\begin{array}{rl}\log p_{\varphi }(y|x_t)& \approx \log p_{\varphi }(y|x_t){|}_{x_t=\mu }+(x_t-\mu ){\nabla }_{x_t}\log p_{\varphi }(y|x_t){|}_{x_t=\mu }\\ & =(x_t-\mu )g+C_1\end{array}$$

• $g={\nabla }_{x_t}\log p_{\varphi }(y|x_t){|}_{x_t=\mu }$
• $C_1\mathrm{is}\mathrm{a}constant.$
  $$\Downarrow$$
  $$\begin{array}{rl}\log (p_{\theta }(x_t|x_{t+1})p_{\varphi }(y|x_t))& =\log p_{\theta }(x_t|x_{t+1})+\log p_{\varphi }(y|x_t)\\ & \approx -\frac{1}{2}(x_t-\mu {)}^T{\Sigma }^{-1}(x_t-\mu )+(x_t-\mu )g+C_2\\ & \stackrel{展开，然后利用二次型的标准形式和对称矩阵性质}{\to }-\frac{1}{2}(x_t-\mu -\Sigma g{)}^T{\Sigma }^{-1}(x_t-\mu -\Sigma g)+\frac{1}{2}{g}^T\Sigma g+C_2\\ & \stackrel{常数项}{\to }-\frac{1}{2}(x_t-\mu -\Sigma g{)}^T{\Sigma }^{-1}(x_t-\mu -\Sigma g)+C_3\\ & \stackrel{转换成高斯分布的对数似然函数}{\to }\boxed{\log p(z)+C_4,z\sim \mathcal{N}(\mu +\Sigma g,\Sigma )}\end{array}$$

在实际应用中，由于直接从这个条件概率分布中采样通常是不可行的，我们需要使用一些近似方法，比如重参数化技巧（reparameterization trick）或者蒙特卡洛方法（Monte Carlo methods），来从这个分布中生成样本
$${\Sigma }_{\theta }\left(X_t,t\right)=exp(vlog{\beta }_t+(1-v)log{\tilde{\beta }}_t)$$
${\tilde{\beta }}_t=\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{\beta }_t$

Classifier 算法

Classifier guided diffusion sampling(stochastic)

与DDPM不同：算出$\mu ，\Sigma$后对均值进行一个偏移g
$$x_{t-1}\leftarrow \mathrm{sample}\mathrm{from}\mathcal{N}(\mu +s\Sigma {\nabla }_{x_t}\log p_{\varphi }(y|x_t),\Sigma )$$

• ${\nabla }_{x_t}\log P_{\Phi }(y|x_t,t)$:当前分类器以$x_t$作为输入时，其输出的概率向量对$x_t$自动微分的一个梯度
• s：缩放量，可以理解为条件执导的程度
  

Classifier guided DDIM sampling

使用score-based conditoning trick
$$\begin{gathered} {\nabla }_{x_t}\log p_{\theta }(x_t)=-\frac{1}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_{\theta }(x_t) \\ \Downarrow \\ \begin{array}{rl}{\nabla }_{x_t}\log (p_{\theta }(x_t)p_{\varphi }(y|x_t))& ={\nabla }_{x_t}\log p_{\theta }(x_t)+{\nabla }_{x_t}\log p_{\varphi }(y|x_t)\\ & =-\frac{1}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_{\theta }(x_t)+{\nabla }_{x_t}\log p_{\varphi }(y|x_t)\end{array} \end{gathered}$$
可以定义DDIM中新的噪音量：
$$\widehat{ϵ}(x_t):={ϵ}_{\theta }(x_t)-\sqrt{1-{\bar{\alpha }}_t}{\nabla }_{x_t}\log p_{\varphi }(y|x_t)$$

Conditional Sampling for DDIM

代码详解

classifier

分类器的作用通常是对输入的数据进行分类或者识别，根据模型的训练，可以对输入的图像或数据进行分类预测。

#\guided_diffusion\script_util.py
def create_classifier(
    image_size,  # 图像大小
    classifier_use_fp16,  # 是否使用 FP16
    classifier_width,  # 分类器宽度
    classifier_depth,  # 分类器深度
    classifier_attention_resolutions,  # 分类器注意力分辨率
    classifier_use_scale_shift_norm,  # 是否使用尺度平移归一化
    classifier_resblock_updown,  # 是否使用上下采样的残差块
    classifier_pool,  # 池化方式
):
    # 根据图像大小选择通道倍增率
    if image_size == 512:
        channel_mult = (0.5, 1, 1, 2, 2, 4, 4)
    elif image_size == 256:
        channel_mult = (1, 1, 2, 2, 4, 4)
    elif image_size == 128:
        channel_mult = (1, 1, 2, 3, 4)
    elif image_size == 64:
        channel_mult = (1, 2, 3, 4)
    else:
        raise ValueError(f"不支持的图像大小: {image_size}")

    # 解析分类器注意力分辨率
    attention_ds = []
    for res in classifier_attention_resolutions.split(","):
        attention_ds.append(image_size // int(res))

    # 创建 EncoderUNetModel 模型
    return EncoderUNetModel(
        image_size=image_size,
        in_channels=3,  # 输入通道数为 3
        model_channels=classifier_width,  # 模型通道数
        out_channels=1000,  # 输出通道数为 1000
        num_res_blocks=classifier_depth,  # 残差块数量
        attention_resolutions=tuple(attention_ds),  # 注意力分辨率
        channel_mult=channel_mult,  # 通道倍增率
        use_fp16=classifier_use_fp16,  # 是否使用 FP16
        num_head_channels=64,  # 注意力头通道数
        use_scale_shift_norm=classifier_use_scale_shift_norm,  # 是否使用尺度平移归一化
        resblock_updown=classifier_resblock_updown,  # 是否使用上下采样的残差块
        pool=classifier_pool,  # 池化方式
    )

cond_fn和model_fn

用于计算分类器的梯度，这个梯度随后可以用于引导扩散模型的采样过程

• cond_fn返回的是$s\times {▽}_{X_t}log{p}_{\varphi }\left(y|X_t\right)$，s 是 args.classifier_scale
• model_fn 函数用于生成样本，它接受输入样本 x、时间步 t 和标签 y 作为参数。它使用扩散模型 model 生成样本，并根据 args.class_cond 参数决定是否使用标签 y 进行条件引导。

#scripts/classifier_sample.py
def main():
    args = create_argparser().parse_args()  # 解析命令行参数

    dist_util.setup_dist()  # 设置分布式计算环境
    logger.configure()  # 配置日志记录

    logger.log("creating model and diffusion...")  # 记录日志
    model, diffusion = create_model_and_diffusion(
        **args_to_dict(args, model_and_diffusion_defaults().keys())
    )  # 创建模型和扩散对象
    model.load_state_dict(
        dist_util.load_state_dict(args.model_path, map_location="cpu")
    )  # 加载模型参数
    model.to(dist_util.dev())  # 将模型移动到指定设备
    if args.use_fp16:
        model.convert_to_fp16()  # 将模型转换为FP16精度
    model.eval()  # 设置模型为评估模式

    logger.log("loading classifier...")  # 记录日志
    classifier = create_classifier(**args_to_dict(args, classifier_defaults().keys()))  # 创建分类器
    classifier.load_state_dict(
        dist_util.load_state_dict(args.classifier_path, map_location="cpu")
    )  # 加载分类器参数
    classifier.to(dist_util.dev())  # 将分类器移动到指定设备
    if args.classifier_use_fp16:
        classifier.convert_to_fp16()  # 将分类器转换为FP16精度
    classifier.eval()  # 设置分类器为评估模式
def cond_fn(x, t, y=None):
#x 是当前的样本数据，t 是时间步，y 是可选的标签。
    assert y is not None
    # 启用梯度计算
    with th.enable_grad():
        # 将输入 x 转换为可求导的张量
        #从 x 创建一个新的变量 x_in，这个变量是 x 的一个副本，但是它的 requires_grad 属性被设置为 True，这意味着在反向传播中需要计算它的梯度。
        x_in = x.detach().requires_grad_(True)
        
        # 使用分类器模型计算输入 x 在时间步 t 的 logits
        logits = classifier(x_in, t)
        
        # 计算 logits 的对数 softmax 概率，这给出了每个类别的对数概率。
        log_probs = F.log_softmax(logits, dim=-1)
        
        # 获取对应于标签 y 的对数概率
        selected = log_probs[range(len(logits)), y.view(-1)]
        
        # 计算对数概率关于输入 x_in 的梯度
        return th.autograd.grad(selected.sum(), x_in)[0] * args.classifier_scale

def model_fn(x, t, y=None):
    assert y is not None
    # 使用扩散模型生成样本
    return model(x, t, y if args.class_cond else None)

logger.log("sampling...")
all_images = []
all_labels = []
# 循环生成样本，直到生成指定数量的样本
while len(all_images) * args.batch_size < args.num_samples:
    model_kwargs = {}
    # 随机生成一批标签
    classes = th.randint(
        low=0, high=NUM_CLASSES, size=(args.batch_size,), device=dist_util.dev()
    )
    # 将标签添加到模型参数中
    model_kwargs["y"] = classes
    # 选择采样方法，使用 p_sample_loop 或 ddim_sample_loop
    sample_fn = (
        diffusion.p_sample_loop if not args.use_ddim else diffusion.ddim_sample_loop
    )
    # 使用采样方法生成样本
    sample = sample_fn(
        model_fn,
        (args.batch_size, 3, args.image_size, args.image_size),
        clip_denoised=args.clip_denoised,
        model_kwargs=model_kwargs,
        cond_fn=cond_fn,
        device=dist_util.dev(),
    )
    # 将样本归一化到 0-255 之间
    sample = ((sample + 1) * 127.5).clamp(0, 255).to(th.uint8)
    # 将样本的维度调整为 (batch_size, height, width, channel)
    sample = sample.permute(0, 2, 3, 1)
    # 将样本转换为连续内存
    sample = sample.contiguous()
    # 将生成的样本和标签添加到列表中
    gathered_samples = [th.zeros_like(sample) for _ in range(dist.get_world_size())]
    dist.all_gather(gathered_samples, sample)  # 收集生成的图像
    all_images.extend([sample.cpu().numpy() for sample in gathered_samples])  # 扩展图像列表
    gathered_labels = [th.zeros_like(classes) for _ in range(dist.get_world_size())]
    dist.all_gather(gathered_labels, classes)  # 收集标签
    all_labels.extend([labels.cpu().numpy() for labels in gathered_labels])  # 扩展标签列表
    logger.log(f"created {len(all_images) * args.batch_size} samples")  # 记录生成的样本数量

p_sample

• p_sample 函数：

  • 该函数从模型在给定时间步 t 上采样 x_{t-1}。
  • 它首先使用 p_mean_variance 函数计算 x_{t-1} 的均值和方差。
    然后，它生成一个与 x 形状相同的噪声 noise。
  • 如果 cond_fn 不为空，则使用 condition_mean 函数根据条件修正均值 out[“mean”]。
  • 最后，它使用均值、方差和噪声计算 x_{t-1} 的样本，并返回一个包含样本和 x_0 预测的字典。

• p_sample_loop 函数：

  • 该函数从模型中生成样本，并返回一个不可微分的样本批次。
  • 使用 p_sample_loop_progressive 函数生成样本，并返回最后一个时间步的样本。

• p_sample_loop_progressive 函数：

  • 该函数从模型中生成样本，并从扩散的每个时间步生成中间样本。
  • 它首先根据 shape 生成一个噪声 img。
  • 然后，它从 self.num_timesteps 到 0 反向遍历时间步 i。在每个时间步，它使用
  • p_sample 函数生成一个样本，并将该样本作为字典返回。它还将当前样本 img 更新为新生成的样本。

与之前DDPM不同的部分：mean要修正

#**p_sample中的cond_fn在condition_mean处起作用**
    if cond_fn is not None:
        out["mean"] = self.condition_mean(  
        # 对 mean 进行修正
            cond_fn, out, x, t, model_kwargs=model_kwargs
        )

def p_sample(
    self,
    model,
    x,
    t,
    clip_denoised=True,
    denoised_fn=None,
    cond_fn=None,
    model_kwargs=None,
):
    """
    从模型在给定时间步上采样 x_{t-1}。

    :param model: 要从中采样的模型。
    :param x: 当前时间步 x_{t-1} 的张量。
    :param t: t 的值，从第一个扩散步骤的 0 开始。
    :param clip_denoised: 如果为 True，则将 x_start 预测裁剪到 [-1, 1]。
    :param denoised_fn: 如果不为 None，则是一个应用于 x_start 预测的函数，
                        在用于采样之前。
    :param cond_fn: 如果不为 None，则是一个梯度函数，其作用类似于模型。
    :param model_kwargs: 如果不为 None，则是一个包含要传递给模型的额外关键字参数的字典。
                        这可用于条件化。
    :return: 包含以下键的字典：
             - 'sample': 从模型中随机采样的样本。
             - 'pred_xstart': x_0 的预测值。
    """
    out = self.p_mean_variance(
        # 基于 xt 计算一些 x_{t-1} 的变量
        model,
        x,
        t,
        clip_denoised=clip_denoised,
        denoised_fn=denoised_fn,
        model_kwargs=model_kwargs,
    )
    noise = th.randn_like(x)
    nonzero_mask = (
        (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
    )  # 当 t == 0 时没有噪声
    if cond_fn is not None:
        out["mean"] = self.condition_mean(  
        # 对 mean 进行修正
            cond_fn, out, x, t, model_kwargs=model_kwargs
        )
    sample = out["mean"] + nonzero_mask * th.exp(0.5 * out["log_variance"]) * noise
    #重参数化
    return {"sample": sample, "pred_xstart": out["pred_xstart"]}


def p_sample_loop(
    self,
    model,
    shape,
    noise=None,
    clip_denoised=True,
    denoised_fn=None,
    cond_fn=None,
    model_kwargs=None,
    device=None,
    progress=False,
):
    """
    从模型中生成样本。

    :param model: 模型模块。
    :param shape: 样本的形状，(N, C, H, W)。
    :param noise: 如果指定，则来自编码器的噪声用于采样。
                  应该与 `shape` 形状相同。
    :param clip_denoised: 如果为 True，则将 x_start 预测裁剪到 [-1, 1]。
    :param denoised_fn: 如果不为 None，则是一个应用于 x_start 预测的函数，
                        在用于采样之前。
    :param cond_fn: 如果不为 None，则是一个梯度函数，其作用类似于模型。
    :param model_kwargs: 如果不为 None，则是一个包含要传递给模型的额外关键字参数的字典。
                        这可用于条件化。
    :param device: 如果指定，则用于创建样本的设备。
                   如果未指定，则使用模型参数的设备。
    :param progress: 如果为 True，则显示 tqdm 进度条。
    :return: 一个不可微分的样本批次。
    """
    final = None
    for sample in self.p_sample_loop_progressive(
        model,
        shape,
        noise=noise,
        clip_denoised=clip_denoised,
        denoised_fn=denoised_fn,
        cond_fn=cond_fn,
        model_kwargs=model_kwargs,
        device=device,
        progress=progress,
    ):
        final = sample
    return final["sample"]


def p_sample_loop_progressive(
    self,
    model,
    shape,
    noise=None,
    clip_denoised=True,
    denoised_fn=None,
    cond_fn=None,
    model_kwargs=None,
    device=None,
    progress=False,
):
    """
    从模型中生成样本，并从扩散的每个时间步生成中间样本。

    参数与 p_sample_loop() 相同。
    返回一个字典生成器，其中每个字典都是 p_sample() 的返回值。
    """
    if device is None:
        device = next(model.parameters()).device
    assert isinstance(shape, (tuple, list))
    if noise is not None:
        img = noise
    else:
        img = th.randn(*shape, device=device)
    indices = list(range(self.num_timesteps))[::-1]

    if progress:
        # 延迟导入，这样我们就不依赖于 tqdm。
        from tqdm.auto import tqdm

        indices = tqdm(indices)

    for i in indices:
        t = th.tensor([i] * shape[0], device=device)
        with th.no_grad():
            out = self.p_sample(
                model,
                img,
                t,
                clip_denoised=clip_denoised,
                denoised_fn=denoised_fn,
                cond_fn=cond_fn,
                model_kwargs=model_kwargs,
            )
            yield out
            img = out["sample"]

ddim_sample

与DDIM不同的部分，DDIM详细推导指路这篇博客：https://blog.csdn.net/hwjokcq/article/details/138873313?spm=1001.2014.3001.5501

if cond_fn is not None:
    out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs)

#guided_diffusion\gaussian_diffusion.py
def ddim_sample(
        self,
        model,
        x,
        t,
        clip_denoised=True,
        denoised_fn=None,
        cond_fn=None,
        model_kwargs=None,
        eta=0.0,
    ):
        """
        Sample x_{t-1} from the model using DDIM.

        Same usage as p_sample().
        """
        # Calculate mean and variance of the prediction
        out = self.p_mean_variance(
            model,
            x,
            t,
            clip_denoised=clip_denoised,
            denoised_fn=denoised_fn,
            model_kwargs=model_kwargs,
        )
        
        # Condition the score if a condition function is provided
        if cond_fn is not None:
            out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs)

        # Predict epsilon based on x_start or x_prev prediction
        eps = self._predict_eps_from_xstart(x, t, out["pred_xstart"])
        
        # Calculate alpha_bar and sigma
        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
        alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape)
        sigma = (
            eta
            * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar))
            * th.sqrt(1 - alpha_bar / alpha_bar_prev)
        )
        
        # Generate noise and predict the mean
        noise = th.randn_like(x)
        mean_pred = (
            out["pred_xstart"] * th.sqrt(alpha_bar_prev)
            + th.sqrt(1 - alpha_bar_prev - sigma ** 2) * eps
        )
        
        # Apply noise only for t != 0
        nonzero_mask = (
            (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
        )
        sample = mean_pred + nonzero_mask * sigma * noise
        
        return {"sample": sample, "pred_xstart": out["pred_xstart"]}

    def ddim_reverse_sample(
        self,
        model,
        x,
        t,
        clip_denoised=True,
        denoised_fn=None,
        model_kwargs=None,
        eta=0.0,
    ):
        """
        Sample x_{t+1} from the model using DDIM reverse ODE.
        """
        # Ensure eta is 0 for deterministic path
        assert eta == 0.0, "Reverse ODE only for deterministic path"
        
        # Calculate mean and variance of the prediction
        out = self.p_mean_variance(
            model,
            x,
            t,
            clip_denoised=clip_denoised,
            denoised_fn=denoised_fn,
            model_kwargs=model_kwargs,
        )
        
        # Predict epsilon based on the model outputs
        eps = (
            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x.shape) * x
            - out["pred_xstart"]
        ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x.shape)
        
        alpha_bar_next = _extract_into_tensor(self.alphas_cumprod_next, t, x.shape)

        # Calculate the mean prediction for the reverse ODE
        mean_pred = (
            out["pred_xstart"] * th.sqrt(alpha_bar_next)
            + th.sqrt(1 - alpha_bar_next) * eps
        )

        return {"sample": mean_pred, "pred_xstart": out["pred_xstart"]}

    def ddim_sample_loop(
        self,
        model,
        shape,
        noise=None,
        clip_denoised=True,
        denoised_fn=None,
        cond_fn=None,
        model_kwargs=None,
        device=None,
        progress=False,
        eta=0.0,
    ):
        """
        Generate samples from the model using DDIM.

        Same usage as p_sample_loop().
        """
        final = None
        for sample in self.ddim_sample_loop_progressive(
            model,
            shape,
            noise=noise,
            clip_denoised=clip_denoised,
            denoised_fn=denoised_fn,
            cond_fn=cond_fn,
            model_kwargs=model_kwargs,
            device=device,
            progress=progress,
            eta=eta,
        ):
            final = sample
        return final["sample"]

    def ddim_sample_loop_progressive(
        self,
        model,
        shape,
        noise=None,
        clip_denoised=True,
        denoised_fn=None,
        cond_fn=None,
        model_kwargs=None,
        device=None,
        progress=False,
        eta=0.0,
    ):
        """
        Use DDIM to sample from the model and yield intermediate samples from
        each timestep of DDIM.

        Same usage as p_sample_loop_progressive().
        """
        if device is None:
            device = next(model.parameters()).device
        assert isinstance(shape, (tuple, list))
        if noise is not None:
            img = noise
        else:
            img = th.randn(*shape, device=device)
        indices = list(range(self.num_timesteps))[::-1]

        if progress:
            # Lazy import so that we don't depend on tqdm.
            from tqdm.auto import tqdm

            indices = tqdm(indices)

        for i in indices:
            t = th.tensor([i] * shape[0], device=device)
            with th.no_grad():
                out = self.ddim_sample(
                    model,
                    img,
                    t,
                    clip_denoised=clip_denoised,
                    denoised_fn=denoised_fn,
                    cond_fn=cond_fn,
                    model_kwargs=model_kwargs,
                    eta=eta,
                )
                yield out
                img = out["sample"]

#guided_diffusion\gaussian_diffusion.py
def condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
        """
        计算在模型的评分函数受到 cond_fn 条件约束时，p_mean_variance 输出会是什么。

        有关 cond_fn 的详细信息，请参阅 condition_mean()。

        与 condition_mean() 不同，这里使用的是 Song 等人（2020年）的条件策略。

        Parameters:
        - cond_fn: 条件函数，用于对模型的评分函数进行条件约束。
        - p_mean_var: 模型输出的平均值和方差。
        - x: 输入数据。
        - t: 时间步长。
        - model_kwargs: 模型参数（可选）。

        Returns:
        - out: 更新后的 p_mean_var 输出。

        """
        # 计算 alpha_bar
        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)

        # 预测 epsilon
        eps = self._predict_eps_from_xstart(x, t, p_mean_var["pred_xstart"])
        # 根据 Song 等人的策略，对 epsilon 进行调整
        eps = eps - (1 - alpha_bar).sqrt() * cond_fn(
            x, self._scale_timesteps(t), **model_kwargs
        )

        # 复制 p_mean_var
        out = p_mean_var.copy()
        # 通过调整后的 epsilon 预测 x_start
        out["pred_xstart"] = self._predict_xstart_from_eps(x, t, eps)
        # 计算 posterior 的均值和方差
        out["mean"], _, _ = self.q_posterior_mean_variance(
            x_start=out["pred_xstart"], x_t=x, t=t
        )
        
        return out