别再死磕公式了！用PyTorch从零实现一个DDPM图像生成模型（附完整代码）-编程实验室

从零构建DDPM图像生成模型：PyTorch实战指南与代码解析

1. 为什么选择DDPM？

在生成对抗网络(GAN)和变分自编码器(VAE)主导生成模型的今天，Denoising Diffusion Probabilistic Models (DDPM)以其独特的训练稳定性和高质量的生成效果异军突起。与GAN容易出现的模式坍塌问题不同，DDPM通过渐进式的加噪和去噪过程，实现了更加可控的图像生成。

DDPM的核心优势：

训练过程稳定，不易出现模式坍塌
生成质量高，尤其在细节保留方面表现优异
理论框架严谨，数学基础扎实
可解释性强，每个步骤都有明确的物理意义

2. DDPM基础架构解析

2.1 前向加噪过程

DDPM的前向过程是一个固定的马尔可夫链，逐步将数据转化为高斯噪声。这个过程可以用以下公式描述：

def forward_process(x0, t, beta): """ x0: 原始图像 t: 时间步 beta: 噪声调度参数 """ sqrt_alpha = torch.sqrt(1 - beta) noise = torch.randn_like(x0) xt = sqrt_alpha * x0 + (1 - sqrt_alpha) * noise return xt

关键参数说明：

beta: 控制噪声添加速率的超参数
alpha: 1 - beta，表示保留原始信息的比例
t: 时间步，决定当前加噪的程度

2.2 反向去噪过程

反向过程是DDPM的核心，通过学习一个神经网络来逐步去除噪声：

class ReverseProcess(nn.Module): def __init__(self): super().__init__() self.model = UNet() # 通常使用U-Net结构 def forward(self, xt, t): predicted_noise = self.model(xt, t) return predicted_noise

3. 完整PyTorch实现

3.1 噪声调度器

合理的噪声调度对模型性能至关重要：

class NoiseScheduler: def __init__(self, timesteps=1000, beta_start=1e-4, beta_end=0.02): self.timesteps = timesteps self.betas = torch.linspace(beta_start, beta_end, timesteps) self.alphas = 1. - self.betas self.alpha_bars = torch.cumprod(self.alphas, dim=0) def sample_noise_level(self, n): t = torch.randint(0, self.timesteps, (n,)) sqrt_alpha_bar = torch.sqrt(self.alpha_bars[t]) sqrt_one_minus_alpha_bar = torch.sqrt(1. - self.alpha_bars[t]) return t, sqrt_alpha_bar, sqrt_one_minus_alpha_bar

3.2 U-Net架构设计

U-Net是DDPM常用的骨干网络：

class UNetBlock(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) def forward(self, x, t): h = self.conv1(x) time_emb = self.time_mlp(t) h = h + time_emb[:, :, None, None] return self.conv2(h)

3.3 训练循环实现

def train_step(model, scheduler, x0, optimizer): model.train() optimizer.zero_grad() # 采样噪声级别 t, sqrt_alpha_bar, sqrt_one_minus_alpha_bar = scheduler.sample_noise_level(x0.shape[0]) # 前向加噪 noise = torch.randn_like(x0) xt = sqrt_alpha_bar[:, None, None, None] * x0 + sqrt_one_minus_alpha_bar[:, None, None, None] * noise # 预测噪声 predicted_noise = model(xt, t) # 计算损失 loss = F.mse_loss(predicted_noise, noise) loss.backward() optimizer.step() return loss.item()

4. 采样生成过程

训练完成后，我们可以通过反向过程生成新图像：

@torch.no_grad() def sample(model, scheduler, img_size, batch_size=8, channels=3): model.eval() xt = torch.randn((batch_size, channels, img_size, img_size)) for t in reversed(range(scheduler.timesteps)): # 预测噪声 noise_pred = model(xt, torch.full((batch_size,), t)) # 计算去噪后的图像 alpha_t = scheduler.alphas[t] alpha_bar_t = scheduler.alpha_bars[t] beta_t = scheduler.betas[t] if t > 0: noise = torch.randn_like(xt) else: noise = torch.zeros_like(xt) xt = (1 / torch.sqrt(alpha_t)) * (xt - ((1 - alpha_t) / torch.sqrt(1 - alpha_bar_t)) * noise_pred) + torch.sqrt(beta_t) * noise return xt

5. 实战技巧与优化

5.1 学习率调度

def get_lr_scheduler(optimizer, warmup_steps=5000): def lr_lambda(step): if step < warmup_steps: return float(step) / float(max(1, warmup_steps)) return 1.0 return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)

5.2 混合精度训练

scaler = torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): predicted_noise = model(xt, t) loss = F.mse_loss(predicted_noise, noise) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()

5.3 可视化训练过程

def plot_samples(samples, nrow=8): grid = torchvision.utils.make_grid(samples, nrow=nrow) plt.imshow(grid.permute(1, 2, 0).cpu().numpy()) plt.axis('off') plt.show()

6. 常见问题与解决方案

问题1：生成图像模糊

增加模型容量
延长训练时间
调整噪声调度参数

问题2：训练不稳定

使用梯度裁剪
调整学习率
增加批量大小

问题3：生成多样性不足

检查损失函数是否过于激进
确保噪声添加足够随机
验证模型没有过拟合

7. 进阶优化方向

对于希望进一步提升模型性能的开发者，可以考虑以下方向：

条件生成：通过添加类别标签或文本描述实现可控生成
加速采样：研究减少采样步数的方法，如DDIM
多尺度架构：在不同分辨率层次上处理图像
混合模型：结合GAN等其他生成模型优势

# 条件生成示例 class ConditionalDDPM(nn.Module): def __init__(self, num_classes): super().__init__() self.class_embed = nn.Embedding(num_classes, 128) self.model = UNet() def forward(self, xt, t, class_labels): class_emb = self.class_embed(class_labels) return self.model(xt, t, class_emb)

8. 实际应用案例

DDPM已经在多个领域展现出强大潜力：

艺术创作：生成独特风格的艺术作品
数据增强：为小样本学习任务生成训练数据
图像修复：补全缺失或损坏的图像区域
医学影像：生成合成医学图像用于研究

# 图像修复示例 def inpainting(model, corrupted_img, mask): xt = corrupted_img for t in reversed(range(timesteps)): # 混合已知区域和生成区域 known_part = mask * corrupted_img + (1 - mask) * xt predicted = model(known_part, t) xt = known_part * mask + predicted * (1 - mask) return xt

9. 性能评估指标

评估生成模型质量常用指标：

指标	描述	实现方式
FID	衡量生成图像与真实图像的分布距离	使用预训练网络提取特征
IS	评估生成图像的多样性和质量	分类网络预测结果分析
PSNR	峰值信噪比，评估重建质量	计算均方误差的对数

# FID计算示例 def calculate_fid(real_imgs, fake_imgs): inception = torchvision.models.inception_v3(pretrained=True) real_features = inception(real_imgs)[0].squeeze() fake_features = inception(fake_imgs)[0].squeeze() mu_real, sigma_real = real_features.mean(0), torch.cov(real_features.T) mu_fake, sigma_fake = fake_features.mean(0), torch.cov(fake_features.T) diff = mu_real - mu_fake cov_mean = (sigma_real @ sigma_fake).sqrt() return diff.dot(diff) + torch.trace(sigma_real + sigma_fake - 2 * cov_mean)