别再死记硬背公式了！用Python和PyTorch手把手拆解Diffusion Model的前向加噪与反向去噪

用Python和PyTorch实战拆解Diffusion Model：从加噪到去噪的完整实现

Diffusion Model近年来在图像生成领域掀起了一场革命，但许多开发者面对其复杂的数学公式时常常感到无从下手。本文将通过代码实战的方式，带你一步步实现Diffusion Model的核心过程——前向加噪与反向去噪。我们将使用PyTorch构建一个极简但完整的Diffusion模型，通过可视化每个步骤的结果，让你直观理解这个强大模型的工作原理。

1. 理解Diffusion Model的基本原理

Diffusion Model的核心思想是通过逐步加噪和去噪的过程来学习数据分布。这个过程可以分为两个阶段：

• 前向过程（Forward Process）：将一张清晰图片逐步添加噪声，最终变成纯噪声
• 反向过程（Reverse Process）：学习如何从噪声中逐步恢复原始图片

提示：Diffusion Model的独特之处在于它不直接生成图像，而是学习如何"去噪"，这使得训练过程更加稳定。

让我们用一个简单的表格对比这两个过程：

2. 环境准备与基础配置

在开始编码前，我们需要设置好开发环境。以下是所需的Python库：

import torch import torch.nn as nn import torch.optim as optim import numpy as np import matplotlib.pyplot as plt from torchvision import datasets, transforms from torch.utils.data import DataLoader

我们将使用MNIST数据集作为示例，因为它相对简单，便于我们观察Diffusion过程的效果。数据加载的代码如下：

transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,)) ]) train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform) train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)

3. 实现前向加噪过程

前向过程的核心公式是：

Xₜ = √αₜ Xₜ₋₁ + √(1-αₜ) εₜ

其中εₜ是从标准正态分布中采样的噪声。我们需要定义一个噪声调度（noise schedule），决定αₜ如何随时间变化。

def linear_beta_schedule(timesteps, start=0.0001, end=0.02): return torch.linspace(start, end, timesteps) timesteps = 200 betas = linear_beta_schedule(timesteps) alphas = 1. - betas alphas_cumprod = torch.cumprod(alphas, axis=0)

接下来实现前向扩散过程的关键函数：

def forward_diffusion_sample(x_0, t, device="cpu"): noise = torch.randn_like(x_0) sqrt_alphas_cumprod_t = torch.sqrt(alphas_cumprod[t])[:, None, None, None] sqrt_one_minus_alphas_cumprod_t = torch.sqrt(1. - alphas_cumprod[t])[:, None, None, None] return sqrt_alphas_cumprod_t.to(device) * x_0.to(device) + sqrt_one_minus_alphas_cumprod_t.to(device) * noise.to(device), noise.to(device)

让我们可视化一下前向过程的效果：

def show_forward_process(sample_img, num_steps=5): plt.figure(figsize=(15, 3)) plt.subplot(1, num_steps+1, 1) plt.imshow(sample_img.squeeze(), cmap='gray') plt.title("Original") for i in range(1, num_steps+1): t = torch.tensor([i * (timesteps // num_steps)]) noisy_img, _ = forward_diffusion_sample(sample_img, t) plt.subplot(1, num_steps+1, i+1) plt.imshow(noisy_img.squeeze().cpu().numpy(), cmap='gray') plt.title(f"Step {t.item()}") plt.show()

4. 构建去噪U-Net模型

反向过程需要一个能够预测噪声的神经网络。我们将实现一个简化的U-Net结构：

class Block(nn.Module): def __init__(self, in_ch, out_ch, time_emb_dim): super().__init__() self.time_mlp = nn.Linear(time_emb_dim, out_ch) self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1) self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1) self.relu = nn.ReLU() def forward(self, x, t): h = self.relu(self.conv1(x)) time_emb = self.relu(self.time_mlp(t)) time_emb = time_emb[(..., ) + (None, ) * 2] h = h + time_emb return self.relu(self.conv2(h)) class SimpleUNet(nn.Module): def __init__(self): super().__init__() self.time_mlp = nn.Sequential( nn.Linear(1, 32), nn.ReLU(), nn.Linear(32, 32) ) self.down1 = Block(1, 32, 32) self.down2 = Block(32, 64, 32) self.up1 = Block(64, 32, 32) self.up2 = Block(32, 1, 32) def forward(self, x, t): t = self.time_mlp(t) x1 = self.down1(x, t) x2 = self.down2(x1, t) x = self.up1(x2, t) return self.up2(x, t)

5. 训练Diffusion模型

有了模型结构后，我们需要定义训练过程。Diffusion Model的训练目标是让网络能够预测添加到图像中的噪声。

device = "cuda" if torch.cuda.is_available() else "cpu" model = SimpleUNet().to(device) optimizer = optim.Adam(model.parameters(), lr=1e-3) epochs = 20 for epoch in range(epochs): for step, (batch, _) in enumerate(train_loader): batch = batch.to(device) t = torch.randint(0, timesteps, (batch.shape[0],), device=device).long() noisy_images, noise = forward_diffusion_sample(batch, t, device) predicted_noise = model(noisy_images, t.unsqueeze(1).float()) loss = nn.functional.mse_loss(predicted_noise, noise) optimizer.zero_grad() loss.backward() optimizer.step() print(f"Epoch {epoch} | Loss: {loss.item():.4f}")

6. 实现反向去噪过程

训练完成后，我们可以使用模型从随机噪声开始逐步去噪生成图像。反向过程的每一步都需要：

1. 预测当前图像的噪声
2. 使用预测的噪声计算更"干净"的图像

@torch.no_grad() def sample_timestep(x, t): betas_t = betas[t].to(device) sqrt_one_minus_alphas_cumprod_t = torch.sqrt(1. - alphas_cumprod[t]).to(device) sqrt_recip_alphas_t = torch.sqrt(1.0 / alphas[t]).to(device) model_mean = sqrt_recip_alphas_t * (x - betas_t * model(x, t.unsqueeze(1).float()) / sqrt_one_minus_alphas_cumprod_t) if t == 0: return model_mean else: posterior_variance_t = betas[t].to(device) noise = torch.randn_like(x) return model_mean + torch.sqrt(posterior_variance_t) * noise @torch.no_grad() def sample_plot_image(): img_size = 28 img = torch.randn((1, 1, img_size, img_size), device=device) plt.figure(figsize=(15, 3)) num_images = 10 stepsize = timesteps // num_images for i in range(0, timesteps)[::-1]: img = sample_timestep(img, torch.tensor([i], device=device)) if i % stepsize == 0: plt.subplot(1, num_images, i // stepsize + 1) plt.imshow(img.squeeze().cpu().numpy(), cmap='gray') plt.title(f"{i}") plt.show()

7. 优化与改进方向

虽然我们的简易实现已经能够展示Diffusion Model的核心原理，但在实际应用中还有很大的优化空间：

• 更复杂的噪声调度：尝试cosine调度而非线性调度
• 更强大的网络结构：使用标准的U-Net架构而非简化版
• 条件生成：加入类别信息实现可控生成
• 加速采样：研究DDIM等快速采样方法

以下是一个改进后的噪声调度实现示例：

def cosine_beta_schedule(timesteps, s=0.008): 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.clamp(betas, 0.0001, 0.9999)

在实际项目中，我发现调整噪声调度对生成质量影响很大。线性调度简单直接，但cosine调度通常在图像质量上有更好表现。另一个关键点是模型容量——对于更复杂的数据集如CIFAR-10或CelebA，需要更深的网络和更多的训练时间才能获得理想结果。