DDPM实战用Python从头实现一个简易扩散模型附完整代码如果你对生成式AI背后的技术感到好奇尤其是那些能凭空创造出逼真图像的模型那么“扩散模型”这个词你一定不陌生。从DALL·E 2到Stable Diffusion这些令人惊叹的AI艺术工具背后Denoising Diffusion Probabilistic Models (DDPM) 扮演了核心角色。很多教程会带你深入数学的海洋从变分推断讲到朗之万动力学但对于想亲手实现一个的开发者来说最迫切的问题往往是“这些公式到底怎么变成能跑的代码”这篇文章就是为你准备的。我们将暂时放下繁复的数学推导聚焦于实战。我将假设你熟悉Python和PyTorch的基本操作并带着你像搭积木一样一步步构建一个简易但完整的DDPM模型。我们的目标不是复现一个SOTA的图像生成器而是亲手实现前向加噪、噪声预测网络和反向去噪采样这三个核心过程并可视化每一个关键步骤。当你看到随机噪声如何一步步被“雕刻”成有意义的图像轮廓时你对DDPM的理解将不再停留在纸面。1. 环境准备与核心概念代码化在开始写模型之前我们需要确保环境就绪并将DDPM中那些关键但略显抽象的概念转化为具体的变量和函数。首先准备好你的Python环境。你需要安装PyTorch用于构建和训练神经网络、Matplotlib用于可视化以及一些标准的数据处理库。如果你使用CUDA显卡请确保安装了对应版本的PyTorch以获得加速。pip install torch torchvision matplotlib numpy接下来我们首先要代码化DDPM最核心的调度器——噪声调度。它决定了在扩散过程的每一步我们要添加多少噪声。在DDPM的原始论文中使用的是线性增长的方差调度但实践中更常用的是余弦调度因为它能带来更平滑的噪声变化和更好的生成质量。我们先实现一个简单的线性调度作为起点。import torch import numpy as np import matplotlib.pyplot as plt def linear_beta_schedule(timesteps, start0.0001, end0.02): 生成线性的beta调度表。 beta_t 决定了第t步添加的噪声量。 return torch.linspace(start, end, timesteps) # 设置总时间步T T 1000 betas linear_beta_schedule(T) # 根据betas计算其他关键派生量 alphas 1. - betas alphas_cumprod torch.cumprod(alphas, dim0) # \bar{\alpha}_t alphas_cumprod_prev torch.cat([torch.tensor([1.0]), alphas_cumprod[:-1]]) sqrt_alphas_cumprod torch.sqrt(alphas_cumprod) # 用于前向过程 sqrt_one_minus_alphas_cumprod torch.sqrt(1. - alphas_cumprod) # 用于前向过程注意alphas_cumprod记作 \bar{\alpha}_t是DDPM中的一个核心变量。在前向过程中任意时刻t的图像x_t可以直接由原始图像x_0加噪得到x_t sqrt(\bar{\alpha}_t) * x_0 sqrt(1-\bar{\alpha}_t) * 噪声。这个重参数化技巧让我们无需逐步模拟1000步就能采样任意中间状态极大提升了训练效率。为了直观感受调度我们可以快速绘制一下关键量的变化曲线plt.figure(figsize(12, 4)) plt.subplot(1, 3, 1) plt.plot(betas.numpy()) plt.title(Beta Schedule (Noise Level)) plt.xlabel(Timestep t) plt.ylabel(Beta_t) plt.subplot(1, 3, 2) plt.plot(alphas_cumprod.numpy()) plt.title(Alpha Cumprod) plt.xlabel(Timestep t) plt.ylabel(\\bar{\\alpha}_t) plt.subplot(1, 3, 3) plt.plot(sqrt_one_minus_alphas_cumprod.numpy()) plt.title(Noise Coefficient) plt.xlabel(Timestep t) plt.ylabel(sqrt(1 - \\bar{\\alpha}_t)) plt.tight_layout() plt.show()你会看到\bar{\alpha}_t从1单调递减至接近0而噪声系数则从0单调递增至接近1。这完美对应了“拆楼”的比喻随着步骤t增加原始图像的信息x_0的系数越来越少随机噪声的成分越来越多直到最后x_T几乎变成纯高斯噪声。2. 实现前向扩散与数据加载理解了调度我们就可以实现前向扩散过程了。这个过程是确定的不包含任何可学习参数。我们将编写一个函数它接收一批原始图像x_0和一个时间步t然后根据上述公式直接返回加噪后的x_t以及所使用的随机噪声noise。这个noise至关重要因为在训练时我们的模型目标就是预测它。def q_sample(x_start, t, noiseNone): 前向扩散过程根据给定的时间步t为输入x_start添加噪声。 使用重参数化技巧一步到位得到x_t。 if noise is None: noise torch.randn_like(x_start) # 获取对应时间步的系数 sqrt_alphas_cumprod_t sqrt_alphas_cumprod[t].view(-1, 1, 1, 1) sqrt_one_minus_alphas_cumprod_t sqrt_one_minus_alphas_cumprod[t].view(-1, 1, 1, 1) # 计算 x_t x_t sqrt_alphas_cumprod_t * x_start sqrt_one_minus_alphas_cumprod_t * noise return x_t, noise为了测试我们的模型需要一个数据集。这里我们使用经典的Fashion-MNIST它复杂度适中训练速度快足以演示DDPM的完整流程。我们将构建一个简单的数据加载管道。from torchvision import datasets, transforms from torch.utils.data import DataLoader # 定义图像预处理转换为Tensor并归一化到[-1, 1]区间 transform transforms.Compose([ transforms.ToTensor(), transforms.Lambda(lambda x: (x * 2) - 1) # 从[0,1]映射到[-1,1] ]) # 加载数据集 train_dataset datasets.FashionMNIST(root./data, trainTrue, downloadTrue, transformtransform) train_loader DataLoader(train_dataset, batch_size128, shuffleTrue, num_workers2, pin_memoryTrue) # 可视化一些原始图像 def show_images(images, title): images (images 1) / 2 # 反归一化到[0,1]以便显示 grid torchvision.utils.make_grid(images[:16], nrow4) plt.imshow(grid.permute(1, 2, 0).numpy()) plt.title(title) plt.axis(off) plt.show() sample_batch, _ next(iter(train_loader)) show_images(sample_batch, Original Fashion-MNIST Images)现在让我们将前向扩散过程可视化。我们取一张图片展示它在不同扩散步数t下的状态。def visualize_forward_diffusion(image, timesteps_to_show[0, 50, 150, 300, 600, 999]): 可视化一张图片在前向扩散过程中的变化。 fig, axes plt.subplots(1, len(timesteps_to_show), figsize(15, 3)) image image.unsqueeze(0) # 增加批次维度 for i, t in enumerate(timesteps_to_show): t_tensor torch.tensor([t]) x_t, _ q_sample(image, t_tensor) img_to_show (x_t.squeeze() 1) / 2 # 反归一化 axes[i].imshow(img_to_show.numpy().squeeze(), cmapgray) axes[i].set_title(ft {t}) axes[i].axis(off) plt.suptitle(Forward Diffusion Process (Image Gradually Becomes Noise)) plt.show() # 选取第一张图片进行可视化 single_image sample_batch[0] visualize_forward_diffusion(single_image)你会清晰地看到图像从清晰逐渐变得模糊最终在t999时几乎变成纯粹的随机噪声。这个过程就是“拆楼”它将数据分布逐渐转化为一个简单的高斯分布。3. 构建U-Net噪声预测模型现在进入核心部分构建一个能够学习“逆过程”的神经网络。在DDPM中这个网络的目标不是直接预测去噪后的图像x_{t-1}而是预测在前向过程中添加到x_t上的噪声ε。这是一个非常巧妙的设计它将一个复杂的图像生成问题转化为了一个相对简单的噪声预测问题。我们将使用一个基于U-Net架构的模型。U-Net最初为图像分割设计其编码器-解码器结构带有跳跃连接非常适合捕捉图像的局部和全局特征在图像生成任务中表现优异。我们的简易U-Net将包含以下几个部分下采样块编码器 由卷积层、组归一化GroupNorm和SiLU激活函数组成逐步降低空间分辨率增加通道数以提取高层语义特征。上采样块解码器 由转置卷积或插值上采样层、组归一化和激活函数组成逐步恢复空间分辨率减少通道数以重建图像细节。跳跃连接 将编码器对应层的特征图与解码器的特征图在通道维度上拼接帮助解码器恢复更精确的空间信息。时间步嵌入 将时间步t通过正弦位置编码或MLP转换为向量并注入到网络的每一层中让模型知道当前正在处理哪个去噪步骤。以下是模型的核心代码实现import torch.nn as nn import torch.nn.functional as F class SinusoidalPositionEmbeddings(nn.Module): 将标量时间步t转换为高维向量表示。 def __init__(self, dim): super().__init__() self.dim dim def forward(self, time): device time.device half_dim self.dim // 2 embeddings np.log(10000) / (half_dim - 1) embeddings torch.exp(torch.arange(half_dim, devicedevice) * -embeddings) embeddings time[:, None] * embeddings[None, :] embeddings torch.cat((embeddings.sin(), embeddings.cos()), dim-1) return embeddings class Block(nn.Module): 基础卷积块Conv - GroupNorm - Activation。 def __init__(self, in_ch, out_ch, time_emb_dimNone): super().__init__() self.conv nn.Conv2d(in_ch, out_ch, 3, padding1) self.norm nn.GroupNorm(8, out_ch) # 8个组是常用设置 self.act nn.SiLU() # 如果提供time_emb_dim则创建一个线性层来处理时间嵌入 if time_emb_dim is not None: self.time_mlp nn.Linear(time_emb_dim, out_ch) else: self.time_mlp None def forward(self, x, t_embNone): h self.conv(x) h self.norm(h) # 如果存在时间嵌入将其加到特征图上 if self.time_mlp is not None and t_emb is not None: time_emb self.time_mlp(t_emb)[:, :, None, None] # 调整形状为 [B, C, 1, 1] h h time_emb h self.act(h) return h class SimpleUNet(nn.Module): 一个简化的U-Net用于预测噪声。 def __init__(self, in_channels1, out_channels1, time_emb_dim32): super().__init__() self.time_mlp nn.Sequential( SinusoidalPositionEmbeddings(time_emb_dim), nn.Linear(time_emb_dim, time_emb_dim), nn.SiLU(), nn.Linear(time_emb_dim, time_emb_dim) ) # 下采样路径 self.down1 Block(in_channels, 64, time_emb_dim) self.down2 Block(64, 128, time_emb_dim) self.down3 Block(128, 256, time_emb_dim) self.bottleneck Block(256, 512, time_emb_dim) # 上采样路径 self.up3 Block(512 256, 256, time_emb_dim) # 跳跃连接增加了通道数 self.up2 Block(256 128, 128, time_emb_dim) self.up1 Block(128 64, 64, time_emb_dim) # 最终输出层预测与输入图像同尺寸的噪声 self.final_conv nn.Conv2d(64, out_channels, kernel_size1) def forward(self, x, timesteps): # 1. 处理时间嵌入 t_emb self.time_mlp(timesteps) # 2. 编码器下采样 d1 self.down1(x, t_emb) d2 self.down2(F.max_pool2d(d1, 2), t_emb) d3 self.down3(F.max_pool2d(d2, 2), t_emb) bottleneck self.bottleneck(F.max_pool2d(d3, 2), t_emb) # 3. 解码器上采样 跳跃连接 u3 self.up3(torch.cat([F.interpolate(bottleneck, scale_factor2), d3], dim1), t_emb) u2 self.up2(torch.cat([F.interpolate(u3, scale_factor2), d2], dim1), t_emb) u1 self.up1(torch.cat([F.interpolate(u2, scale_factor2), d1], dim1), t_emb) # 4. 最终输出 return self.final_conv(u1)这个模型虽然结构简单但包含了DDPM噪声预测网络的所有关键要素。我们可以快速实例化它并检查输入输出的维度是否匹配。# 实例化模型并测试前向传播 device torch.device(cuda if torch.cuda.is_available() else cpu) model SimpleUNet(in_channels1, out_channels1).to(device) print(fModel initialized on {device}) # 创建一个随机输入模拟一个批次的图像和随机时间步 test_batch_size 4 test_image torch.randn(test_batch_size, 1, 28, 28).to(device) # Fashion-MNIST是28x28灰度图 test_timesteps torch.randint(0, T, (test_batch_size,)).to(device) # 前向传播 with torch.no_grad(): predicted_noise model(test_image, test_timesteps) print(fInput shape: {test_image.shape}) print(fPredicted noise shape: {predicted_noise.shape}) # 应与输入形状一致4. 训练循环与损失函数有了模型和前向扩散过程我们就可以定义训练流程了。DDPM的训练目标非常简洁最小化模型预测的噪声与在前向过程中实际添加的噪声之间的均方误差MSE Loss。训练循环的每一步可以分解为从数据加载器中获取一个批次的真实图像x_0。为批次中的每一张图像随机采样一个时间步t均匀分布在[0, T-1]。使用q_sample函数根据采样的t为x_0添加噪声得到x_t和真实噪声noise。将x_t和t输入噪声预测模型得到预测的噪声predicted_noise。计算predicted_noise与noise之间的均方误差作为损失。反向传播更新模型参数。以下是训练循环的代码实现def train_one_epoch(model, loader, optimizer, device, epoch): model.train() total_loss 0 for batch_idx, (clean_images, _) in enumerate(loader): clean_images clean_images.to(device) batch_size clean_images.shape[0] # 1. 随机采样时间步t t torch.randint(0, T, (batch_size,), devicedevice).long() # 2. 前向扩散添加噪声 noisy_images, true_noise q_sample(clean_images, t) # 3. 模型预测噪声 predicted_noise model(noisy_images, t) # 4. 计算简单的均方误差损失 loss F.mse_loss(predicted_noise, true_noise) # 5. 反向传播与优化 optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() if batch_idx % 100 0: print(fEpoch {epoch} | Batch {batch_idx:04d} | Loss: {loss.item():.6f}) avg_loss total_loss / len(loader) print(f Epoch {epoch} Average Loss: {avg_loss:.6f}) return avg_loss # 初始化优化器 optimizer torch.optim.AdamW(model.parameters(), lr1e-4) # 训练若干个Epoch num_epochs 20 for epoch in range(1, num_epochs 1): avg_loss train_one_epoch(model, train_loader, optimizer, device, epoch) # 这里可以添加模型保存和验证逻辑在训练初期损失会下降得比较快因为模型在学习噪声的基本模式。随着训练进行损失下降会变缓模型在学习更精细的、与图像内容相关的噪声结构。即使只训练10-20个epoch在Fashion-MNIST这样相对简单的数据集上你也能观察到模型开始具备一定的去噪能力。5. 反向采样生成图像模型训练完成后最激动人心的部分来了从纯噪声中生成图像。反向采样是一个迭代过程从x_T ~ N(0, I)开始逐步应用训练好的模型进行去噪直到得到x_0。根据DDPM的推导采样公式如下x_{t-1} (1 / sqrt(alpha_t)) * (x_t - ((1-alpha_t) / sqrt(1-alpha_bar_t)) * model(x_t, t)) sigma_t * z其中当t 1时z ~ N(0, I)为随机噪声当t 1时z 0。sigma_t是方差通常可以设置为sqrt(beta_t)。让我们实现这个采样循环并观察图像是如何一步步“浮现”出来的。torch.no_grad() def p_sample(model, x, t, t_index): 单步反向去噪采样。 x: 当前时刻的带噪图像 x_t t: 当前时间步张量 t_index: 当前时间步的索引整数 # 获取对应时间步的系数 betas_t betas[t].view(-1, 1, 1, 1).to(x.device) sqrt_one_minus_alphas_cumprod_t sqrt_one_minus_alphas_cumprod[t].view(-1, 1, 1, 1).to(x.device) sqrt_recip_alphas_t torch.sqrt(1.0 / (1 - betas_t)) # 模型预测噪声 pred_noise model(x, t) # 计算 x_{t-1} 的均值 model_mean sqrt_recip_alphas_t * (x - betas_t * pred_noise / sqrt_one_minus_alphas_cumprod_t) if t_index 0: # 最后一步不再添加随机噪声 return model_mean else: # 添加随机噪声 posterior_variance_t betas[t].to(x.device) noise torch.randn_like(x) return model_mean torch.sqrt(posterior_variance_t).view(-1, 1, 1, 1) * noise torch.no_grad() def p_sample_loop(model, shape, device): 完整的反向采样循环。 shape: 要生成图像的形状 (batch_size, channels, height, width) # 从纯噪声开始 img torch.randn(shape, devicedevice) imgs [] # 可选保存中间过程 for i in reversed(range(0, T)): t torch.full((shape[0],), i, devicedevice, dtypetorch.long) img p_sample(model, img, t, i) # 每100步保存一次中间结果用于可视化 if i % 100 0: imgs.append(img.cpu()) imgs.append(img.cpu()) # 保存最终结果 return imgs torch.no_grad() def generate_and_visualize(model, device, num_images16): 生成一批图像并可视化采样过程。 model.eval() # 生成图像 sample_shape (num_images, 1, 28, 28) generated_imgs_list p_sample_loop(model, sample_shape, device) # 可视化最终结果 final_imgs generated_imgs_list[-1] show_images(final_imgs, fGenerated Images (Final Step)) # 可视化采样过程选择其中一张图像的演变 fig, axes plt.subplots(1, len(generated_imgs_list), figsize(15, 3)) for idx, img_step in enumerate(generated_imgs_list): # 取批次中的第一张图 single_img img_step[0].unsqueeze(0) img_to_show (single_img.squeeze() 1) / 2 axes[idx].imshow(img_to_show.numpy().squeeze(), cmapgray) step (len(generated_imgs_list)-1-idx) * 100 if idx len(generated_imgs_list)-1 else 0 axes[idx].set_title(fStep {step}) axes[idx].axis(off) plt.suptitle(Reverse Denoising Process (From Noise to Image)) plt.tight_layout() plt.show() # 调用函数生成图像 generate_and_visualize(model, device)运行这段代码后你将看到两幅图。第一幅是最终生成的一批Fashion-MNIST风格的图像可能有些模糊或扭曲这取决于模型训练的程度。第二幅图则动态展示了其中一张图像从纯噪声t1000开始每经过100步去噪后的状态。你会清晰地看到一个模糊的轮廓逐渐变得清晰最终形成一个可识别的衣物或鞋类图像。这个过程就是“建楼”模型根据学到的知识将无序的噪声“重建”为有结构的数据。6. 关键技巧、优化与扩展方向一个基础的DDPM已经实现了。但在实际研究和应用中还有许多技巧可以显著提升效果和效率。这里分享几个关键点你可以基于此进行实验1. 改进的噪声调度我们之前使用了简单的线性调度。尝试换成余弦调度它能提供更平滑的噪声变化通常在视觉质量上表现更好。def cosine_beta_schedule(timesteps, s0.008): Cosine schedule as proposed in https://arxiv.org/abs/2102.09672 steps timesteps 1 x torch.linspace(0, timesteps, steps) alphas_cumprod torch.cos(((x / timesteps) s) / (1 s) * torch.pi * 0.5) ** 2 alphas_cumprod alphas_cumprod / alphas_cumprod[0] betas 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1]) return torch.clip(betas, 0.0001, 0.9999) # 使用余弦调度 betas_cosine cosine_beta_schedule(T) # ... 后续用 betas_cosine 重新计算 alphas, alphas_cumprod 等2. 更强大的网络架构我们的SimpleUNet非常基础。现代扩散模型通常使用更多的通道数和残差块。注意力机制特别是自注意力用于建模图像不同区域之间的长程依赖这对生成连贯的物体至关重要。条件注入例如将类别标签或文本描述编码后注入模型实现可控生成。3. 加速采样算法DDPM需要迭代T步通常1000步才能生成一张图速度很慢。研究者提出了很多加速方法其核心思想是减少采样步数同时通过更复杂的更新公式来弥补步数减少带来的误差。例如DDIM(Denoising Diffusion Implicit Models) 就是一种经典的加速采样方法它可以将采样步数减少到50甚至20步而质量损失很小。4. 损失函数的变体我们使用了最简单的噪声MSE损失。一些工作探索了其他损失如基于感知的损失或者对噪声预测误差进行重新加权例如给中间时间步更高的权重以提升生成图像的清晰度和细节。5. 从Fashion-MNIST到更复杂的数据集一旦你在简单数据集上跑通了整个流程挑战就变成了如何让模型在更复杂、更高分辨率的数据集如CIFAR-10, CelebA甚至ImageNet上工作。这通常需要更深的U-Net和注意力层。在图像像素空间直接训练非常消耗内存因此通常会使用潜在扩散模型 (LDM)即先用一个VAE或VQ-VAE将图像压缩到低维潜在空间然后在潜在空间中进行扩散和生成最后再用解码器还原到像素空间。这是Stable Diffusion采用的核心技术。实现这个简易DDPM的最大收获是亲手验证了“拆楼-建楼”这一直观思想如何通过代码变为现实。最初的生成结果可能不完美但每一次训练、每一次采样都是对概率分布如何被逐步塑造的深刻体验。我建议你尝试修改超参数如学习率、T的大小、更换噪声调度、或者在U-Net中增加一层注意力观察这些变化如何影响生成图像的质量和多样性。真正的理解始于你动手改变代码并看到结果变化的那一刻。