Diffusion model in Pytorch

---

前言

Diffusion model 是机器学习中的一个比较新的生成式模型，且有望和GAN一较高下，在这里记录一些我的学习和理解。

---

一、代码参考位置

Github: denoising-diffusion-pytorch
网址链接：https://github.com/lucidrains/denoising-diffusion-pytorch

二、代码关联关系解析

1. 调用的packages之间的关系

三个类扮演不同的角色实现去噪扩散模型的训练和图像生成过程。

1. Unet的作用
   U-Net 架构的神经网络模型，用于图像分割任务，在去噪扩散模型中，U-Net被用作生成模型(条件模型)，每一个扩散步骤中预测噪声，并逐渐将带有噪声的图像转换为清晰的图像。

2. GaussianDiffusion的作用
   实现去噪的核心算法，负责定义扩散过程，和反扩散过程。
   训练阶段：将图像逐渐添加噪声，训练神经网络逆转噪声过程
   生成阶段：扩散模型从纯噪声开始，逐步去噪生成最终的图像

3. Trainer的作用
   Trainer 类负责设置和执行模型的训练过程。它管理数据加载、批量处理、优化器设置、模型保存、性能监控等训练任务。

可以说Gaussian Diffusion是一个更广泛的生成模型框架，而U-Net可以作为这个框架中的一个组件。

2. 详细代码

from denoising_diffusion_pytorch import Unet, GaussianDiffusion, Trainer

#Unet,GuassianDiffusion,Trainer的作用

model = Unet(
    dim=64, #网络的基础通道数
    dim_mults=(1, 2, 4, 8), #每个下采样阶段的通道数是前一个阶段的倍数
    flash_attn=True
)

diffusion = GaussianDiffusion(
    model, #传入的Unet模型，预测每一步去噪所需的噪声
    image_size=128, #表示生成的图像大小
    timesteps=1000, #整个扩散过程所分成的步骤
    sampling_timesteps=250    # 表示在生成图像时，采样过程使用250个时间步骤（使用更少的步骤可以加快生成速度
)

trainer = Trainer(
    diffusion, #传入的GuassianDiffusion对象
    'path/to/your/images', #训练图像数据的文件路径
    train_batch_size=32, #训练时每个批次的图像数量
    train_lr=8e-5, #学习率
    train_num_steps=700000, #训练的总步数
    gradient_accumulate_every=2,    #每两个批次累积一次梯度更新
    ema_decay=0.995,         #指数移动平均（EMA）的衰减率，用于平滑模型参数
    amp=True,                # 使用自动混合精度训练，可以提高训练速度并节省内存
    calculate_fid=True              # Frechet Inception Distance（FID），这是评估生成图像质量的指标。
)

trainer.train()

# 获取生成的图像结果
sampled_images = diffusion.sample(batch_size=4)
print(sampled_images.shape) # 输出结果 (4, 3, 128, 128)

---

总结

例如：以上就是今天要讲的内容，本文简单介绍了 Unet, GaussianDiffusion, Trainer在 Diffusion model 中的相互制约关系。