一文读懂Denoising Diffusion模型架构：从UNet到采样器的完整解析
项目地址: https://gitcode.com/gh_mirrors/de/denoising-diffusion-pytorch 项目概述
Denoising Diffusion Probabilistic Model（去噪扩散概率模型）是一种能够生成高质量图像的深度学习模型。本项目GitHub 加速计划 / de / denoising-diffusion-pytorch提供了该模型的PyTorch实现，通过简洁的代码结构和模块化设计，使研究者和开发者能够快速理解和应用扩散模型。
模型架构总览
扩散模型主要由两个核心部分组成：UNet网络和扩散过程。UNet负责学习从含噪图像中预测噪声，而扩散过程则控制噪声的添加和去除。整个模型的工作流程如下：
- 前向扩散：将高斯噪声逐步添加到原始图像中
- 反向扩散：利用UNet网络逐步从含噪图像中去除噪声，最终生成清晰图像
UNet网络详解
基本结构
项目中的UNet实现位于denoising_diffusion_pytorch/denoising_diffusion_pytorch.py文件中，采用了经典的编码器-解码器架构，并加入了注意力机制以捕捉长距离依赖关系。
UNet的主要组件包括：
- 输入卷积层：将输入图像映射到特征空间
- 下采样模块：通过卷积和池化操作提取图像特征
- 上采样模块：通过转置卷积和跳跃连接恢复图像细节
- 注意力模块：在网络深层加入自注意力机制
核心代码解析
UNet类的初始化函数定义了网络的整体结构：
class Unet(Module):
def __init__(
self,
dim,
init_dim = None,
out_dim = None,
dim_mults = (1, 2, 4, 8),
channels = 3,
self_condition = False,
learned_variance = False,
learned_sinusoidal_cond = False,
random_fourier_features = False,
learned_sinusoidal_dim = 16,
sinusoidal_pos_emb_theta = 10000,
dropout = 0.,
attn_dim_head = 32,
attn_heads = 4,
full_attn = None, # defaults to full attention only for inner most layer
flash_attn = False
):
super().__init__()
# 网络初始化代码
其中，dim_mults参数控制了不同层级的特征维度，通过下采样和上采样操作实现特征的提取和恢复。
时间嵌入
扩散模型的关键创新之一是将时间信息嵌入到网络中。项目中实现了两种时间嵌入方式：
- 正弦位置嵌入：基于三角函数的固定嵌入
- 随机或学习的正弦嵌入：可学习的参数化嵌入
时间嵌入代码位于denoising_diffusion_pytorch/denoising_diffusion_pytorch.py的SinusoidalPosEmb和RandomOrLearnedSinusoidalPosEmb类中。
注意力机制
为了捕捉图像中的长距离依赖关系，UNet在深层加入了注意力机制。项目实现了两种注意力模块：
- 线性注意力：计算效率高，适合高分辨率图像
- 标准注意力：捕捉能力强，计算成本较高
注意力模块的实现位于denoising_diffusion_pytorch/denoising_diffusion_pytorch.py的LinearAttention和Attention类中。
Karras UNet变体
除了标准UNet，项目还提供了Karras等人提出的改进版UNet，位于denoising_diffusion_pytorch/karras_unet.py文件中。该变体引入了以下改进：
- 幅度保持操作：确保网络各层的输出幅度一致
- 强制权重归一化：提高训练稳定性
- 改进的注意力机制：增强特征捕捉能力
Karras UNet的核心改进在于引入了幅度保持的卷积和激活函数，如MPSiLU和PixelNorm，这些改进有助于稳定训练过程并提高生成质量。
扩散过程详解
前向扩散
前向扩散过程将高斯噪声逐步添加到原始图像中，使得图像在扩散步数T后变成完全的高斯噪声。项目中实现了多种噪声调度策略：
- 线性调度：噪声线性增加
- 余弦调度：噪声按照余弦曲线增加
- Sigmoid调度：噪声按照Sigmoid曲线增加
噪声调度的实现位于denoising_diffusion_pytorch/denoising_diffusion_pytorch.py的linear_beta_schedule、cosine_beta_schedule和sigmoid_beta_schedule函数中。
反向扩散与采样
反向扩散过程是利用UNet网络从噪声中逐步恢复图像的过程。项目实现了多种采样方法：
- DDPM采样：原始扩散模型的采样方法，需要较多步数
- DDIM采样：加速采样方法，可显著减少采样步数
采样过程的实现位于denoising_diffusion_pytorch/denoising_diffusion_pytorch.py的GaussianDiffusion类中，主要通过p_sample_loop和ddim_sample方法实现。
连续时间扩散
项目还实现了连续时间扩散模型，位于denoising_diffusion_pytorch/continuous_time_gaussian_diffusion.py文件中。该模型将扩散过程视为连续时间的随机过程，通过求解随机微分方程实现图像生成。
模型训练与目标函数
目标函数
扩散模型的训练目标是最小化预测噪声与真实噪声之间的均方误差。项目中实现了多种目标函数变体：
- 预测噪声：直接预测添加的噪声
- 预测原始图像：直接预测去噪后的图像
- 预测v参数：预测噪声和图像的组合参数
这些目标函数的实现位于denoising_diffusion_pytorch/denoising_diffusion_pytorch.py的model_predictions方法中。
训练技巧
为了提高训练效率和生成质量，项目集成了多种训练技巧：
- 最小SNR损失权重：调整不同扩散阶段的损失权重
- 偏移噪声：在训练中引入偏移噪声，提高多样性
- EMA（指数移动平均）：平滑模型参数，提高稳定性
实际应用示例
使用该项目训练的扩散模型可以生成高质量图像。以下是一个示例生成结果：
要使用该项目训练自己的扩散模型，可以参考项目的README.md文件，其中包含了详细的安装和使用说明。
总结与展望
本项目提供了扩散模型的完整实现，包括标准UNet和Karras UNet变体，以及多种扩散过程和采样方法。通过模块化的设计，研究者可以方便地修改和扩展模型，探索新的扩散模型变体。
未来的改进方向可能包括：
- 进一步优化采样速度
- 提高生成图像的分辨率和质量
- 扩展到视频和3D数据生成
通过本项目，开发者和研究者可以深入理解扩散模型的工作原理，并将其应用于各种生成任务中。
资源与参考
- 项目源代码：gh_mirrors/de/denoising-diffusion-pytorch
- 官方文档：README.md
- 扩散模型原始论文：Denoising Diffusion Probabilistic Models
- Karras UNet论文：Elucidating the Design Space of Diffusion-Based Generative Models
创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
