从零到一!扩散模型核心原理与实战代码全解析

📅2026/7/13 11:02:47 👁️次浏览
从零到一!扩散模型核心原理与实战代码全解析
1. 扩散模型基础概念想象你有一杯清水每次往里面滴一滴墨水经过足够多次后整杯水会变成均匀的黑色。扩散模型的工作原理与此类似只不过我们把滴墨水的过程变成了加噪声。这个看似简单的想法却成为了当前最强大的图像生成技术之一。扩散模型的核心包含两个过程前向过程扩散过程就像往清水里滴墨水逐步将数据如图片变成纯噪声逆向过程相当于时光倒流学习如何从噪声中逐步恢复原始数据我第一次接触这个概念时最惊讶的是为什么这种破坏-重建的方式会比直接生成效果更好后来在实践中发现逐步去噪比直接生成整个图像更容易让模型学习。就像我们修复古画时也是先去除表层污渍再处理深层损伤一步步来反而更有效。2. 前向扩散过程详解2.1 数学原理前向过程可以形式化为马尔可夫链每一步都向数据添加少量高斯噪声。具体来说给定原始图像x₀我们定义一系列逐渐噪声化的版本x₁, x₂,...,x_Tdef forward_diffusion(x0, t, betas): x0: 原始图像 t: 时间步 betas: 噪声调度参数 alphas 1 - betas alpha_bars torch.cumprod(alphas, dim0) noise torch.randn_like(x0) xt torch.sqrt(alpha_bars[t]) * x0 torch.sqrt(1 - alpha_bars[t]) * noise return xt这里betas控制每个时间步添加的噪声量通常从β₁1e-4线性增加到β_T0.02。我曾在实验中调整这个调度参数发现太小的初始beta会导致早期去噪困难而太大的初始beta会让后期保留太多噪声。2.2 噪声调度的艺术选择合适的噪声调度(beta schedule)很关键。常见的有线性调度简单但效果不错余弦调度更平滑的噪声过渡自定义调度针对特定数据集优化# 线性噪声调度 def linear_beta_schedule(timesteps): beta_start 0.0001 beta_end 0.02 return torch.linspace(beta_start, beta_end, timesteps) # 余弦噪声调度来自Improved DDPM def cosine_beta_schedule(timesteps, s0.008): steps timesteps 1 x torch.linspace(0, timesteps, steps) alphas_cumprod torch.cos(((x / timesteps) s) / (1 s) * math.pi * 0.5) ** 2 alphas_cumprod alphas_cumprod / alphas_cumprod[0] betas 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1]) return torch.clip(betas, 0, 0.999)在实际项目中我发现余弦调度通常能生成更清晰的图像特别是当timesteps较大(如1000步)时。而对于快速实验线性调度就足够好了。3. 逆向扩散过程3.1 理论推导逆向过程的目标是学习如何去噪—从x_t预测x_{t-1}。根据论文我们可以推导出q(x_{t-1}|x_t,x_0) N(x_{t-1}; μ̃_t(x_t,x_0), β̃_t I)其中μ̃_t 1/√α_t (x_t - (1-α_t)/√(1-ᾱ_t) ε_t)β̃_t (1-ᾱ_{t-1})/(1-ᾱ_t) β_t这个结果告诉我们如果我们知道原始图像x₀和添加的噪声ε_t就能精确地逆转扩散过程。当然在实际中我们不知道x₀所以需要神经网络来预测。3.2 噪声预测网络我们通常使用U-Net结构的模型来预测噪声class UNet(nn.Module): def __init__(self, dim64, dim_mults(1,2,4,8)): super().__init__() dims [3, *map(lambda m: dim * m, dim_mults)] in_out list(zip(dims[:-1], dims[1:])) self.down_blocks nn.ModuleList([]) self.up_blocks nn.ModuleList([]) # 下采样路径 for ind, (in_dim, out_dim) in enumerate(in_out): self.down_blocks.append( nn.Sequential( nn.Conv2d(in_dim, out_dim, 3, padding1), nn.GroupNorm(8, out_dim), nn.SiLU(), nn.Conv2d(out_dim, out_dim, 3, padding1), nn.GroupNorm(8, out_dim), nn.SiLU(), nn.MaxPool2d(2) if ind len(in_out)-1 else nn.Identity() ) ) # 上采样路径 for ind, (in_dim, out_dim) in enumerate(reversed(in_out[1:])): self.up_blocks.append( nn.Sequential( nn.Conv2d(in_dim*2, out_dim, 3, padding1), nn.GroupNorm(8, out_dim), nn.SiLU(), nn.Conv2d(out_dim, out_dim, 3, padding1), nn.GroupNorm(8, out_dim), nn.SiLU(), nn.Upsample(scale_factor2) if ind len(in_out)-2 else nn.Identity() ) ) self.mid_block nn.Sequential( nn.Conv2d(dims[-1], dims[-1], 3, padding1), nn.GroupNorm(8, dims[-1]), nn.SiLU(), nn.Conv2d(dims[-1], dims[-1], 3, padding1), nn.GroupNorm(8, dims[-1]), nn.SiLU() ) self.final_conv nn.Conv2d(dim, 3, 1) def forward(self, x, t): # 这里简化了时间嵌入的处理 h [] for block in self.down_blocks: x block(x) h.append(x) x self.mid_block(x) for block in self.up_blocks: x torch.cat([x, h.pop()], dim1) x block(x) return self.final_conv(x)这个U-Net有几个关键设计点跳跃连接保留多尺度信息组归一化稳定训练时间步t通过条件注入影响各层4. 训练目标与损失函数扩散模型的训练目标相当简洁—最小化预测噪声和实际噪声的差距def p_losses(denoise_model, x0, t, noiseNone): if noise is None: noise torch.randn_like(x0) xt q_sample(x0, t, noise) # 前向扩散得到xt predicted_noise denoise_model(xt, t) return F.mse_loss(noise, predicted_noise)这个简单的目标函数背后有深厚的理论支撑—它实际上是变分下界(VLB)的简化形式。我在复现DDPM时发现直接预测噪声比预测去噪后的图像更容易训练模型收敛更快。5. 采样生成过程5.1 基础采样算法训练好噪声预测模型后生成新图像的流程如下torch.no_grad() def p_sample(model, x, t, t_index): betas_t extract(betas, t, x.shape) sqrt_one_minus_alphas_cumprod_t extract( sqrt_one_minus_alphas_cumprod, t, x.shape ) sqrt_recip_alphas_t extract(sqrt_recip_alphas, t, x.shape) # 预测噪声 model_mean sqrt_recip_alphas_t * ( x - betas_t * model(x, t) / sqrt_one_minus_alphas_cumprod_t ) if t_index 0: return model_mean else: posterior_variance_t extract(posterior_variance, t, x.shape) noise torch.randn_like(x) return model_mean torch.sqrt(posterior_variance_t) * noise torch.no_grad() def p_sample_loop(model, shape): device next(model.parameters()).device # 从纯噪声开始 img torch.randn(shape, devicedevice) imgs [] for i in tqdm(reversed(range(0, timesteps)), descsampling loop): img p_sample(model, img, torch.full((shape[0],), i, devicedevice, dtypetorch.long), i) imgs.append(img.cpu().numpy()) return imgs5.2 加速采样技巧基础采样需要T步通常T1000这很耗时。后来研究者提出了多种加速方法DDIM将扩散过程重新定义为非马尔可夫过程允许跳步采样知识蒸馏训练学生模型模仿多步教师模型的行为渐进式蒸馏迭代式地减少所需步数我在项目中尝试过DDIM用50步就能获得接近1000步的质量torch.no_grad() def ddim_sample(model, shape, steps50, eta0.0): # 创建时间步序列 times torch.linspace(0, timesteps-1, steps1).long().flip(0) time_pairs list(zip(times[:-1], times[1:])) x torch.randn(shape, devicedevice) x_seq [x.cpu()] for time, time_next in time_pairs: alpha alphas_cumprod[time] alpha_next alphas_cumprod[time_next] # 预测噪声 eps model(x, time) # 计算x0预测 x0_pred (x - torch.sqrt(1. - alpha) * eps) / torch.sqrt(alpha) # 计算方向 dir_xt torch.sqrt(1. - alpha_next) * eps # 更新x x torch.sqrt(alpha_next) * x0_pred dir_xt if eta 0: noise torch.randn_like(x) sigma eta * torch.sqrt((1 - alpha/alpha_next) * (1 - alpha_next)/(1 - alpha)) x sigma * noise x_seq.append(x.cpu()) return x, x_seq6. 实战在MNIST上训练扩散模型6.1 数据准备def get_dataloader(batch_size128): transform Compose([ ToTensor(), Lambda(lambda x: (x - 0.5) * 2) # 将像素值缩放到[-1,1] ]) dataset MNIST(./data, trainTrue, downloadTrue, transformtransform) return DataLoader(dataset, batch_sizebatch_size, shuffleTrue)6.2 模型训练def train(): dataloader get_dataloader() model UNet(dim32, dim_mults(1,2,4)).to(device) optimizer Adam(model.parameters(), lr2e-4) for epoch in range(20): for step, (x, _) in enumerate(dataloader): x x.to(device) # 随机采样时间步 t torch.randint(0, timesteps, (x.shape[0],), devicedevice).long() # 计算损失 loss p_losses(model, x, t) # 反向传播 optimizer.zero_grad() loss.backward() optimizer.step() if step % 100 0: print(fEpoch {epoch} | Step {step} | Loss {loss.item():.4f}) # 每个epoch保存样本 samples p_sample_loop(model, (16, 1, 28, 28)) save_images(samples, fepoch_{epoch}.png)6.3 结果可视化训练约20个epoch后模型应该能生成可辨认的数字。生成过程可视化如下初始噪声 → 中间状态 → 最终生成 ██████ ██▓▓██ ██████ ██░░██ ██▒▒██ ██░░██ ██░░██ ██▒▒██ ██░░██ ██████ ██▓▓██ ██████7. 进阶技巧与优化7.1 条件生成通过添加类别信息可以实现可控生成class ConditionalUNet(UNet): def __init__(self, num_classes10, **kwargs): super().__init__(**kwargs) self.label_emb nn.Embedding(num_classes, kwargs[dim]) def forward(self, x, t, y): emb self.label_emb(y).unsqueeze(-1).unsqueeze(-1) return super().forward(x emb, t)7.2 混合精度训练大幅减少显存使用并加速训练scaler torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): loss p_losses(model, x, t) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()7.3 超参数调优关键超参数及其影响学习率2e-4通常不错太大导致不稳定batch size越大越好但受显存限制模型深度更深的模型容量更大但更难训练timesteps1000足够更多步数提升有限8. 常见问题与解决方案问题1生成图像模糊检查噪声调度是否合理模型容量是否足够解决尝试余弦调度增加模型深度或通道数问题2训练不稳定检查梯度是否爆炸损失是否NaN解决添加梯度裁剪使用更稳定的归一化层问题3生成多样性不足检查模型是否过拟合采样温度是否太低解决增加dropout调整采样时的噪声量9. 扩展应用扩散模型不仅用于图像生成还可应用于图像修复已知部分图像生成缺失部分超分辨率从低分辨率生成高分辨率文本到图像如DALL-E 2、Stable Diffusion3D生成点云或体素生成10. 完整代码示例以下是完整的PyTorch实现import torch import torch.nn as nn import torch.nn.functional as F from torch.optim import Adam from torchvision import transforms from torchvision.datasets import MNIST from torch.utils.data import DataLoader import numpy as np import matplotlib.pyplot as plt from tqdm import tqdm # 参数设置 timesteps 1000 image_size 28 channels 1 batch_size 128 device cuda if torch.cuda.is_available() else cpu # 噪声调度 betas linear_beta_schedule(timesteps) alphas 1. - betas alphas_cumprod torch.cumprod(alphas, axis0) alphas_cumprod_prev F.pad(alphas_cumprod[:-1], (1, 0), value1.0) sqrt_recip_alphas torch.sqrt(1.0 / alphas) sqrt_alphas_cumprod torch.sqrt(alphas_cumprod) sqrt_one_minus_alphas_cumprod torch.sqrt(1. - alphas_cumprod) posterior_variance betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod) # 模型定义 model UNet(dim32, dim_mults(1,2,4)).to(device) optimizer Adam(model.parameters(), lr2e-4) # 训练循环 for epoch in range(20): for step, (x, _) in enumerate(dataloader): x x.to(device) t torch.randint(0, timesteps, (x.shape[0],), devicedevice).long() loss p_losses(model, x, t) optimizer.zero_grad() loss.backward() optimizer.step() # 保存模型和样本 torch.save(model.state_dict(), fddpm_mnist_epoch{epoch}.pt) samples p_sample_loop(model, (16, channels, image_size, image_size)) save_images(samples, fepoch_{epoch}.png)