扩散模型详解：原理与代码实现

扩散模型详解：原理与代码实现 | 极客日志

1. 简介

扩散模型（Diffusion Model）的起源可以追溯到概率图模型和统计物理学领域。它最初的灵感来自于对热扩散和布朗运动等物理现象的研究，这些过程描述了系统如何从一个高能量、不均匀的状态逐步过渡到一个低能量、平衡的状态，后来被引入机器学习和生成模型领域。

扩散模型的早期形式可以追溯到多种研究工作，Sohl-Dickstein 等人（2015 年）在论文《Deep Unsupervised Learning using Nonequilibrium Thermodynamics》中首次将扩散过程引入深度生成模型的框架，奠定了扩散模型的理论基础。

近年来，扩散模型得到了快速发展。尤其是 Ho 等人（2020 年）在论文《Denoising Diffusion Probabilistic Models (DDPM)》中提出了一种高效的扩散模型框架，这一工作开辟了扩散模型的新方向，使其在图像生成、语音合成和其他生成任务中表现优异。DDPM 将扩散过程分为两个阶段，一是前向过程（Forward Process），向数据中逐步添加噪声，直到数据接近高斯分布。二是反向过程（Reverse Process），学习逐步去噪，最终从随机噪声中生成目标数据。

[图：随机噪声经过去噪模块生成猫图像的示意图]

就像是米开朗基罗说的：'塑像就在石头里，我只是把不需要的部分去掉'。

2. 原理

2.1 前向过程

前向过程也就是我们常说的扩散过程，它模拟真实数据逐渐被噪声污染的过程。其做法通常是从高斯分布里面采样一组噪声添加到正常的图片当中，产生有点噪声的图像，然后从高斯分布中再采样一次，再得到更加噪声的图片，以此类推，最后整张图片就看不出来原来是什么东西，也就是整张图片变成了一个接近高斯分布的噪声。做完这个扩散过程以后，就有去噪模块的训练数据了。

前向过程是一个固定的、不可学习的马尔可夫链。从初始数据分布 $x_0$ 开始，逐步向数据中添加高斯噪声，使数据分布逐渐接近标准高斯分布 $\mathcal{N}(0, I)$。公式表示为：

$$ q(x_t | x_{t-1}) = \mathcal{N}(x_t; \sqrt{1 - \beta_t} x_{t-1}, \beta_t I) $$

其中 $x_t$ 是当前时间步的数据，$x_{t-1}$ 是前一时间步的数据，$\beta_t$ 是一个预定义的时间步长参数，表示每一步添加噪声的强度。$I$ 是单位矩阵，表示每个数据维度上的噪声是独立且均匀的。$\mathcal{N}$ 是高斯分布，表示 $x_t$ 的条件分布。

添加噪声的过程是一个马尔可夫链，在每一步中，数据根据高斯分布从 $x_{t-1}$ 生成 $x_t$，噪声的均值为 $0$，方差为 $\beta_t$。前向过程将原始数据逐渐加噪，最终在 $T$ 步后，使其分布接近标准高斯分布 $\mathcal{N}(0, I)$。

2.2 反向过程

反向过程的目的学习如何从完全随机的噪声逐步还原出目标数据。反向过程是需要通过神经网络来学习的，该网络的输入是一张有噪声的图，输出是一张滤掉一点噪声的图像，去噪越做越多，最终就能看到一张清晰的图片。

[图：反向去噪过程示意图]

通常，这个去噪的模型里面实际上是一个噪声预测器（noise predictor），它会预测图片里面的噪声。这个噪声预测器的输入是去噪的图片和噪声现在的严重程度（也就是我们现在进行到去噪的第几个步骤的代号）。它预测在这张图片里面噪声应该长什么样子，再在去噪的图片中减去它预测的噪声，就产生去噪以后的结果，即输出一张噪声少了一点的图。

[图：噪声预测器结构示意]

要训练这样的噪声预测器，要用到之前我们在扩散过程中产生的训练数据。即，扩散过程中产生的一张加完噪声的图片跟现在是第几次加噪声，是网络的输入，而加入的这个噪声就是网络应该要预测的输出。

但有些时候我们不仅想要产生图片，还想产生与我们文字描述一样的图片，对于这样的情况，我们只需要在训练数据中增加对图片的描述，同时在去噪的每一个步骤中让噪声预测器多一个额外的输入，也就是描述的这段文字。

[图：条件扩散模型示意]

反向过程假设数据的逆演化仍然是一个马尔可夫过程，模型需要学习如何从 $x_t$ 预测 $x_{t-1}$，逐步还原出无噪声的原始数据：

$$ p_\theta(x_{t-1} | x_t) = \mathcal{N}(x_{t-1}; \mu_\theta(x_t, t), \Sigma_\theta(x_t, t)) $$

其中 $x_t$ 是当前时间步的带噪声数据，$x_{t-1}$ 是目标时间步的数据（去噪后）。$\mu_\theta$ 是由神经网络预测的均值。$\Sigma_\theta$ 是由神经网络预测的方差（通常固定为常数以简化计算）。通过训练神经网络，使其能够预测每一步中的噪声成分。

扩散模型的训练目标是通过最大化似然估计来优化反向过程。训练过程等价于一个降噪任务，模型学习在给定加噪数据 $x_t$ 的情况下，如何估计当前时间步的噪声，扩散模型的训练损失函数如下：

$$ L = \mathbb{E}{t, x_0, \epsilon} [ || \epsilon - \epsilon\theta(x_t, t) ||^2 ] $$

其中 $\epsilon$ 是前向过程中实际添加的噪声，$\epsilon_\theta$ 是模型预测的噪声。$L_t$ 是当前时间步的损失，表示实际噪声与预测噪声之间的均方误差。通过最小化这个损失，模型学习如何在每一步准确预测噪声，从而能够反向还原数据。

[图：损失函数计算示意]

3. 代码

下面我们以生成动漫人脸图像为目标来训练 Diffusion Model。

动漫人脸数据集下载链接：

https://www.kaggle.com/datasets/b07202024/diffusion/download?datasetVersionNumber=1

本代码遵循典型的 DDPM（Denoising Diffusion Probabilistic Model）框架，整体分为 U-Net 模型（用于去噪）、GaussianDiffusion 类（提供前向扩散和反向采样逻辑）以及数据集和训练器等部分。U-Net 负责在不同尺度下对图像特征进行编码与解码，以预测在每个时间步中加入的噪声；GaussianDiffusion 封装了核心公式与超参数，包括 beta 调度、采样/训练流程及损失函数；Trainer 则管理训练过程，如数据加载、梯度累积、EMA（指数移动平均）等。这种架构将'加噪'和'去噪'分离并封装在模型和调度器中，使得训练和推理流程更加清晰易懂，也能方便地进行扩展或替换不同的网络结构与超参数策略。

以下是完整代码（引自《李宏毅深度学习》）：

import math  # 引入数学相关函数和常量  
import copy  # 提供浅拷贝和深拷贝功能  
from pathlib import Path  # 方便进行路径操作  
from random import random  # 随机数生成函数  
from functools import partial  # 可用于对函数进行部分参数绑定  
from collections import namedtuple  # 提供类似结构体的命名元组  
from multiprocessing import cpu_count  # 获取当前机器的 CPU 核心数  
  
import torch  # 引入 PyTorch 库  
from torch import nn, einsum  # nn 用于神经网络相关模块，einsum 可进行爱因斯坦求和  
import torch.nn.functional as F  # 提供常用函数式神经网络操作  
from torch.utils.data import Dataset, DataLoader  # 数据集基类和数据加载器  
  
from torch.optim import Adam  # Adam 优化器  
  
import torchvision  # PyTorch 视觉工具包  
from torchvision import transforms as T, utils  # 图像预处理 transforms，utils 提供图像显示保存等功能  
  
from einops import rearrange, reduce, repeat  # einops 提供灵活的张量变换函数  
from einops.layers.torch import Rearrange  # einops 在 PyTorch 中的 Layer 实现  
  
from PIL import Image  # Python Image Library，用于图像读写等操作  
 tqdm.auto  tqdm  
 ema_pytorch  EMA  
                              
  
 accelerate  Accelerator  
 matplotlib.pyplot  plt  
 os  
  

torch.backends.cudnn.benchmark =   
  

torch.manual_seed()  
  

 torch.cuda.is_available():  
  torch.cuda.manual_seed()  
  
  
 ():  
      
    scale =  / timesteps  
    beta_start = scale *   
    beta_end = scale *       
     torch.linspace(beta_start, beta_end, timesteps, dtype = torch.float64)  
    
  
  
 ():  
      
    b, *_ = t.shape  
    out = a.gather(-, t)  
     out.reshape(b, *((,) * ((x_shape) - )))  
    
  
  
 ():  
      
     ():  
        .folder = folder        
        .image_size = image_size  
        .paths = [p  p  Path().glob()]  
        
  
        
        
        
        .transform = T.Compose([  
            T.Resize(image_size),  
            T.ToTensor()           
        ])  
  
     ():  
        
         (.paths)  
  
     ():  
        
        path = .paths[index]  
        img = Image.(path)    
         .transform(img)  

 ():  
     x     
    
  
 ():  
     exists(val):  
         val  
     d()  (d)  d  
    
    
  
 ():  
     t  
    
  
 ():  
     :  
         data  dl:  
             data  
    
  
 ():  
     (math.sqrt(num) ** ) == num  
    
    
  
 ():  
    groups = num // divisor  
    remainder = num % divisor  
    arr = [divisor] * groups  
     remainder > :  
        arr.append(remainder)  
     arr  
    
    
  

  
 ():  
     img *  -   
    
  
 ():  
     (t + ) *   
    
  

  
 (nn.Module):  
     ():  
        ().__init__()  
        .fn = fn  
        
  
     ():  
         .fn(x, *args, **kwargs) + x  
        
  
 ():  
     nn.Sequential(  
        nn.Upsample(scale_factor = , mode = ),  
        nn.Conv2d(dim, default(dim_out, dim), , padding = )  
    )  
    
    
    
  
 ():  
     nn.Sequential(  
        Rearrange(, p1 = , p2 = ),  
        nn.Conv2d(dim * , default(dim_out, dim), )  
    )  
    
    
    
    
  
 (nn.Conv2d):  
      
     ():  
        eps =   x.dtype == torch.float32    
        
  
        weight = .weight  
        mean = reduce(weight, , )  
        var = reduce(weight, , partial(torch.var, unbiased = ))  
        
  
        normalized_weight = (weight - mean) * (var + eps).rsqrt()  
        
  
         F.conv2d(x, normalized_weight, .bias, .stride, .padding, .dilation, .groups)  
        
  
 (nn.Module):  
     ():  
        ().__init__()  
        .g = nn.Parameter(torch.ones(, dim, , ))  
        
  
     ():  
        eps =   x.dtype == torch.float32    
        var = torch.var(x, dim = , unbiased = , keepdim = )  
        mean = torch.mean(x, dim = , keepdim = )  
         (x - mean) * (var + eps).rsqrt() * .g  
        
  
 (nn.Module):  
     ():  
        ().__init__()  
        .fn = fn  
        .norm = LayerNorm(dim)  
        
  
     ():  
        x = .norm(x)  
         .fn(x)  
        
  

  
 (nn.Module):  
     ():  
        ().__init__()  
        .dim = dim  
        
  
     ():  
        device = x.device  
        half_dim = .dim //   
        emb = math.log() / (half_dim - )  
        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)  
        
  
        emb = x[:, ] * emb[, :]  
        
  
        emb = torch.cat((emb.sin(), emb.cos()), dim=-)  
        
  
         emb  
        
  
 (nn.Module):  
     following @crowsonkb b -> b d ->  db c -> b c  b (h c) x y -> b h c (x y)b h d n, b h e n -> b h d eb h d e, b h d n -> b h e nb h c (x y) -> b (h c) x yb (h c) x y -> b h c (x y)b h d i, b h d j -> b h i jb h i j, b h d j -> b h i db h (x y) d -> b (h d) x ylinealinealineaunknown beta schedule {beta_schedule}lineabetasalphas_cumprodalphas_cumprod_prevsqrt_alphas_cumprodsqrt_one_minus_alphas_cumprodlog_one_minus_alphas_cumprodsqrt_recip_alphas_cumprodsqrt_recipm1_alphas_cumprodposterior_varianceposterior_log_variance_clippedposterior_mean_coef1posterior_mean_coef2loss_weightsampling loop time stepnoneb ... -> b (...)meanheight  width of image must be {img_size}./faces/faceslinea

accelerate 1.0.1
einops 0.8.0
ema-pytorch 0.7.7
matplotlib 3.5.1
multiprocess 0.70.15
numpy 1.24.4
python 3.8.19
pytorch 2.4.0
pytorch-cuda 12.1
tqdm 4.66.5


下图为模型在完成训练之后生成的动漫人脸图像：

[图：训练完成的动漫人脸生成示例]

从该结果可以看出，模型成功地学习到了二次元人脸的整体特征与色彩分布，生成的人像在发型、五官、配色等方面都有一定的多样性，说明扩散模型在此任务中具备一定的泛化能力。不过图像中仍存在一定程度的模糊、面部细节缺失或扭曲等现象，表明训练规模与网络容量可能还需要进一步优化，以获得更精细、更稳定的生成质量。

总结一下，扩散模型（Diffusion Model）通过在前向过程逐步向图像添加噪声、在反向过程逐步去噪的方式实现图像生成，具有相对稳定的训练过程和良好的生成多样性。它在高分辨率图像生成、条件生成（文本、语音、语义分割等）方面表现不错，且与自回归、GAN 等其他生成方法形成互补。未来发展方向包括更高效的采样策略、更灵活的条件控制、多尺度或多模态的融合，以及在更广泛的数据类型（视频、3D 等）上的应用和研究。

') * 2 * math.pi # 计算随机或可学习的频率 freqs fouriered = torch.cat((freqs.sin(), freqs.cos()), dim = -1) # 拼接正弦和余弦 fouriered = torch.cat((x, fouriered), dim = -1) # 再将原始 x 与正余弦部分合并 return fouriered # 返回包含输入和正余弦编码的结果 # building block modules class Block(nn.Module): def __init__(self, dim, dim_out, groups = 8): super().__init__() self.proj = WeightStandardizedConv2d(dim, dim_out, 3, padding = 1) self.norm = nn.GroupNorm(groups, dim_out) self.act = nn.SiLU() # 一个卷积 -> GroupNorm -> SiLU 激活的基本模块 # 卷积使用 WeightStandardizedConv2d，便于搭配 GroupNorm def forward(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 # 如果从时间嵌入得到 scale_shift，则对特征图进行缩放和偏移 x = self.act(x) return x # 输出经过标准化和激活的张量 class ResnetBlock(nn.Module): def __init__(self, dim, dim_out, *, time_emb_dim = None, groups = 8): super().__init__() # 如果传入了 time_emb_dim，则对时间嵌入进行线性映射得到 scale 和 shift self.mlp = nn.Sequential( nn.SiLU(), nn.Linear(time_emb_dim, dim_out * 2) ) if exists(time_emb_dim) else None self.block1 = Block(dim, dim_out, groups = groups) self.block2 = Block(dim_out, dim_out, groups = groups) self.res_conv = nn.Conv2d(dim, dim_out, 1) if dim != dim_out else nn.Identity() # 如果 dim != dim_out，就用 1x1 卷积在残差分支中对通道数进行调整 def forward(self, x, time_emb = None): scale_shift = None if exists(self.mlp) and exists(time_emb): time_emb = self.mlp(time_emb) time_emb = rearrange(time_emb, '

', x = h, y = w) # reshape 回原始分辨率 return self.to_out(out) # 最后再用 1x1 卷积映射回 dim 维度 # model class Unet(nn.Module): def __init__( self, dim, init_dim = None, out_dim = None, dim_mults=(1, 2, 4, 8), channels = 3, resnet_block_groups = 8, learned_sinusoidal_cond = False, random_fourier_features = False, learned_sinusoidal_dim = 16 ): super().__init__() # determine dimensions self.channels = channels init_dim = default(init_dim, dim) self.init_conv = nn.Conv2d(channels, init_dim, 7, padding = 3) # 输入通道 -> init_dim, 使用 7x7 卷积做初始特征提取 dims = [init_dim, *map(lambda m: dim * m, dim_mults)] in_out = list(zip(dims[:-1], dims[1:])) # 比如 dim=64, dim_mults=(1,2,4,8), 则 dims=[64, 64*1, 64*2, 64*4, 64*8] # in_out 就是 [(64,64),(64,128),(128,256),(256,512)] block_klass = partial(ResnetBlock, groups = resnet_block_groups) # 使用部分函数 partial，将 ResnetBlock 的 groups 参数固定 # time embeddings time_dim = dim * 4 # 时间嵌入的维度，一般设置为 4 倍 base dim self.random_or_learned_sinusoidal_cond = learned_sinusoidal_cond or random_fourier_features if self.random_or_learned_sinusoidal_cond: sinu_pos_emb = RandomOrLearnedSinusoidalPosEmb(learned_sinusoidal_dim, random_fourier_features) fourier_dim = learned_sinusoidal_dim + 1 else: sinu_pos_emb = SinusoidalPosEmb(dim) fourier_dim = dim # 根据需要选择使用随机/可学习的正弦嵌入，或使用经典的正弦嵌入 self.time_mlp = nn.Sequential( sinu_pos_emb, nn.Linear(fourier_dim, time_dim), nn.GELU(), nn.Linear(time_dim, time_dim) ) # 时间嵌入先经过正弦嵌入，然后用两个全连接层（中间激活为 GELU），维度转为 time_dim # layers self.downs = nn.ModuleList([]) self.ups = nn.ModuleList([]) 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.ModuleList([ block_klass(dim_in, 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))), Downsample(dim_in, dim_out) if not is_last else nn.Conv2d(dim_in, dim_out, 3, padding = 1) ])) # down 阶段： # 1) ResnetBlock(dim_in -> dim_in) # 2) 再一个 ResnetBlock(dim_in -> dim_in) # 3) Residual(PreNorm(LinearAttention)) # 4) 如果不是最后层，用 Downsample；否则用 3x3 卷积保持分辨率 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) # 中间层（U-Net 最底部）：ResnetBlock -> 自注意力 -> ResnetBlock for ind, (dim_in, dim_out) in enumerate(reversed(in_out)): is_last = ind == (len(in_out) - 1) # 倒序遍历 in_out，用于 up 阶段 self.ups.append(nn.ModuleList([ block_klass(dim_out + dim_in, dim_out, time_emb_dim = time_dim), block_klass(dim_out + dim_in, dim_out, time_emb_dim = time_dim), Residual(PreNorm(dim_out, LinearAttention(dim_out))), Upsample(dim_out, dim_in) if not is_last else nn.Conv2d(dim_out, dim_in, 3, padding = 1) ])) # up 阶段的逻辑与 down 类似，只是要先拼接 skip connection self.out_dim = default(out_dim, channels) # 最终输出通道数，默认与输入通道一致 self.final_res_block = block_klass(dim * 2, dim, time_emb_dim = time_dim) # 最后一步和初始输入拼接后，再过一个 ResnetBlock self.final_conv = nn.Conv2d(dim, self.out_dim, 1) # 通过 1x1 卷积将维度映射到 out_dim def forward(self, x, time): x = self.init_conv(x) # 初始卷积提取特征 r = x.clone() # 保存初始特征用于最后拼接 t = self.time_mlp(time) # 将时间步 time 通过 time_mlp 得到时间嵌入 t h = [] # 用于保存每层的输出，以便在解码器阶段做 skip connection # downsample for block1, block2, attn, downsample in self.downs: x = block1(x, t) h.append(x) x = block2(x, t) x = attn(x) h.append(x) x = downsample(x) # 依次执行 block1 -> block2 -> attn -> downsample # 并存储中间输出 h # mid x = self.mid_block1(x, t) x = self.mid_attn(x) x = self.mid_block2(x, t) # U-Net 中间层的处理 # upsample for block1, block2, attn, upsample in self.ups: # pop 出下采样时存储的输出，进行 skip connection x = torch.cat((x, h.pop()), dim = 1) x = block1(x, t) x = torch.cat((x, h.pop()), dim = 1) x = block2(x, t) x = attn(x) x = upsample(x) # final x = torch.cat((x, r), dim = 1) # 跟最初的输入特征 r 拼接 x = self.final_res_block(x, t) return self.final_conv(x) # 最终输出一个跟输入维度相匹配的特征图 model = Unet(64) # 实例化一个 U-Net 模型，基本通道数 dim = 64 class GaussianDiffusion(nn.Module): def __init__( self, model, # 传入的 U-Net 等模型，用于预测噪声 *, image_size, # 图像大小（宽和高） timesteps = 1000, # 扩散过程的总时间步数 beta_schedule = '

', maybe_clipped_snr / snr) # 用于加权损失的系数 # auto-normalization of data [0, 1] -> [-1, 1] - can turn off by setting it to be False self.normalize = normalize_to_neg_one_to_one if auto_normalize else identity self.unnormalize = unnormalize_to_zero_to_one if auto_normalize else identity # 根据 auto_normalize 决定是否对数据进行 [-1,1] <-> [0,1] 的转换 def predict_start_from_noise(self, x_t, t, noise): """ 通过 x_t 和噪声，反推 x_0 的预测值 x_0 = 1 / sqrt(alpha_cumprod) * x_t - sqrt(1 / alpha_cumprod - 1) * noise """ return ( extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise ) def predict_noise_from_start(self, x_t, t, x0): """ 通过 x_t 和对 x_0 的预测值，反推噪声的预测值 noise = (1 / sqrt(alpha_cumprod) * x_t - x_0) / sqrt(1 / alpha_cumprod - 1) """ return ( (extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - x0) / \ extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) ) def q_posterior(self, x_start, x_t, t): """ 计算后验分布 q(x_{t-1} | x_t, x_0) 的均值和方差 """ 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) # 后验分布方差的对数（已做 clip） return posterior_mean, posterior_variance, posterior_log_variance_clipped def model_predictions(self, x, t, clip_x_start = False, rederive_pred_noise = False): """ 给定当前噪声图 x 和时间步 t，通过模型预测噪声 pred_noise，并得到对 x_0 的估计 x_start """ model_output = self.model(x, t) # 模型输出，通常是预测噪声 maybe_clip = partial(torch.clamp, min = -1., max = 1.) if clip_x_start else identity # 如果需要对预测出的 x_0 做裁剪，则 partial(torch.clamp)；否则恒等函数 pred_noise = model_output # 这里把模型输出视为噪声预测 x_start = self.predict_start_from_noise(x, t, pred_noise) x_start = maybe_clip(x_start) # 对 x_0 进行 [-1,1] 裁剪（可选） if clip_x_start and rederive_pred_noise: # 如果 x_0 被裁剪，为了更准确，需要重新计算一次噪声 pred_noise = self.predict_noise_from_start(x, t, x_start) return pred_noise, x_start def p_mean_variance(self, x, t, clip_denoised = True): """ 计算从扩散过程中 p(x_{t-1} | x_t) 的均值和方差，用于反向采样 """ noise, x_start = self.model_predictions(x, t) # 模型预测噪声和 x_0 if clip_denoised: x_start.clamp_(-1., 1.) # 默认会把 x_0 的范围裁剪到 [-1,1] model_mean, posterior_variance, posterior_log_variance = self.q_posterior( x_start = x_start, x_t = x, t = t ) # 计算后验分布的均值和方差 # 这里的后验分布相当于 q(x_{t-1}|x_t, x_0) return model_mean, posterior_variance, posterior_log_variance, x_start @torch.no_grad() def p_sample(self, x, t: int): """ 在反向扩散的某一个时间步 t，从 p(x_{t-1} | x_t) 采样 """ b, *_, device = *x.shape, x.device batched_times = torch.full((b,), t, device = x.device, dtype = torch.long) # 构造与批大小相同的时间张量 model_mean, _, model_log_variance, x_start = self.p_mean_variance( x = x, t = batched_times, clip_denoised = True ) # 根据 x_t 计算后验均值和方差 noise = torch.randn_like(x) if t > 0 else 0. # 如果 t > 0 则在采样时加噪声；如果 t=0，则不再加噪声 pred_img = model_mean + (0.5 * model_log_variance).exp() * noise # 采样公式： x_{t-1} = 均值 + 标准差 * 噪声 return pred_img, x_start @torch.no_grad() def p_sample_loop(self, shape, return_all_timesteps = False): """ 从纯噪声开始，逐步反向采样还原图像 """ batch, device = shape[0], self.betas.device # batch 大小，使用存储在 buffer 中的 betas 的设备 img = torch.randn(shape, device = device) # 初始从标准正态分布采样 imgs = [img] # 用于保存采样过程中每个时间步的结果 x_start = None ########################################### ## TODO: plot the sampling process ## ########################################### for t in tqdm(reversed(range(0, self.num_timesteps)), desc = '

', total = self.num_timesteps): # 从 T-1 到 0 逐步反向采样 img, x_start = self.p_sample(img, t) imgs.append(img) ret = img if not return_all_timesteps else torch.stack(imgs, dim = 1) # 如果 return_all_timesteps=True, 返回整个采样序列；否则只返回最终生成的图像 ret = self.unnormalize(ret) # 将图像从 [-1,1] 转回 [0,1] return ret @torch.no_grad() def sample(self, batch_size = 16, return_all_timesteps = False): """ 对外提供的采样接口 """ image_size, channels = self.image_size, self.channels sample_fn = self.p_sample_loop # 默认使用 p_sample_loop 进行逐步采样 return sample_fn((batch_size, channels, image_size, image_size), return_all_timesteps = return_all_timesteps) def q_sample(self, x_start, t, noise=None): """ 前向扩散：从 x_0 得到 x_t 的采样 x_t = sqrt(α_cumprod) * x_0 + sqrt(1-α_cumprod) * noise """ noise = default(noise, lambda: torch.randn_like(x_start)) # 如果不指定噪声，则生成一个和 x_start 形状相同的高斯噪声 return ( extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start + extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise ) @property def loss_fn(self): return F.mse_loss # 训练时使用的损失函数，默认是 MSE def p_losses(self, x_start, t, noise = None): """ 在给定 x_0 以及随机的时间步 t 时，计算训练时的损失 """ b, c, h, w = x_start.shape noise = default(noise, lambda: torch.randn_like(x_start)) # noise sample x = self.q_sample(x_start = x_start, t = t, noise = noise) # 前向扩散，将 x_0 添加噪声到 x_t # predict and take gradient step model_out = self.model(x, t) # 模型对 x_t 进行估计噪声 loss = self.loss_fn(model_out, noise, reduction = '

r' # beta 的调度方式（表示在扩散过程中 beta 的变化），此处设置为线性 model = Unet( dim = channels, dim_mults = dim_mults ) # 实例化一个 U-Net 模型对象，输入的基本通道数为 16， # 会根据 dim_mults 逐步在网络层中增加通道数 diffusion = GaussianDiffusion( model, image_size = IMG_SIZE, timesteps = timesteps, beta_schedule = beta_schedule ) # 将 U-Net 模型封装到 GaussianDiffusion 类中， # 并设置扩散过程中的一些参数（如图像大小、时间步数等）。 # 该类会负责前向扩散（加噪）和反向扩散（去噪）的具体实现。 class Trainer: def __init__( self, diffusion, path, train_batch_size=16, train_lr=1e-3, train_num_steps=10000, gradient_accumulate_every=1, ema_decay=0.995, save_and_sample_every=1000, ): self.accelerator = Accelerator() self.diffusion = diffusion self.batch_size = train_batch_size self.gradient_accumulate_every = gradient_accumulate_every self.train_num_steps = train_num_steps self.ema = EMA(diffusion, decay=ema_decay) self.save_and_sample_every = save_and_sample_every self.path = path self.optimizer = Adam(diffusion.parameters(), lr=train_lr) self.dataset = Dataset(path, image_size=diffusion.image_size) self.dataloader = DataLoader(self.dataset, batch_size=train_batch_size, shuffle=True) self.dataloader = self.accelerator.prepare(self.dataloader) self.diffusion = self.accelerator.prepare(self.diffusion) self.optimizer = self.accelerator.prepare(self.optimizer) def train(self): progress_bar = tqdm(initial=0, total=self.train_num_steps) progress_bar.set_description("Training") for i in range(self.train_num_steps): for img, _ in self.dataloader: loss = self.diffusion(img) self.accelerator.backward(loss / self.gradient_accumulate_every) if i % self.gradient_accumulate_every == 0: self.optimizer.step() self.optimizer.zero_grad() if i % self.save_and_sample_every == 0: self.ema.ema_model.eval() samples = self.ema.ema_model.sample(batch_size=self.batch_size) print(f"Step {i}: Saving samples...") progress_bar.update(1) self.accelerator.end_training() trainer = Trainer( diffusion, path, train_batch_size=batch_size, train_lr=lr, train_num_steps=train_num_steps, gradient_accumulate_every=grad_steps, ema_decay=ema_decay, save_and_sample_every=1000 ) # 实例化一个 Trainer 类来管理训练流程： # - 使用 diffusion 模型进行前向与反向传播 # - 每个 batch 的大小为 16 # - 使用学习率 1e-3 # - 总训练步数为 10000 # - 每个 step 都更新梯度（grad_steps=1） # - EMA 衰减因子为 0.995 # - 每 1000 步保存一次模型并进行一次采样 trainer.train() # 开始训练，Trainer 内部会执行循环读取数据、前向计算、损失反传、优化器更新等流程。运行环境:

扩散模型详解：原理与代码实现

1. 简介

2. 原理

2.1 前向过程

2.2 反向过程

3. 代码

更多推荐文章

相关免费在线工具

扩散模型详解：原理与代码实现

1. 简介

2. 原理

2.1 前向过程

2.2 反向过程

3. 代码

微信扫一扫，关注极客日志

更多推荐文章

相关免费在线工具