实战指南：从零构建PyTorch版Latent Diffusion Models（含DDPM/DDIM/PLMS全流程解析）

1. 环境准备与项目搭建在开始构建Latent Diffusion Models之前我们需要准备好开发环境。这里推荐使用Python 3.8和PyTorch 1.12版本。如果你有GPU设备建议安装CUDA 11.3以上版本以获得更好的训练性能。首先创建一个conda虚拟环境conda create -n ldm python3.8 conda activate ldm然后安装PyTorch和必要的依赖库pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113 pip install numpy tqdm matplotlib tensorboard项目目录结构建议如下latent-diffusion-pytorch/ ├── configs/ # 配置文件 ├── datasets/ # 数据集 ├── models/ # 模型定义 │ ├── autoencoder.py # 变分自编码器 │ ├── diffusion.py # 扩散模型 │ └── unet.py # UNet网络 ├── sampling/ # 采样方法 │ ├── ddpm.py # DDPM采样 │ ├── ddim.py # DDIM采样 │ └── plms.py # PLMS采样 ├── utils/ # 工具函数 ├── train.py # 训练脚本 └── generate.py # 生成脚本2. 变分自编码器(VAE)实现2.1 VAE基础架构变分自编码器是Latent Diffusion Models的核心组件负责将图像压缩到潜在空间。我们实现一个中等压缩比(1/8)的VAEimport torch import torch.nn as nn class VAE(nn.Module): def __init__(self, in_channels3, latent_channels4, base_channels64): super().__init__() # 编码器 self.encoder nn.Sequential( nn.Conv2d(in_channels, base_channels, 4, 2, 1), nn.GroupNorm(32, base_channels), nn.SiLU(), nn.Conv2d(base_channels, base_channels*2, 4, 2, 1), nn.GroupNorm(32, base_channels*2), nn.SiLU(), nn.Conv2d(base_channels*2, base_channels*4, 4, 2, 1), nn.GroupNorm(32, base_channels*4), nn.SiLU(), nn.Conv2d(base_channels*4, latent_channels*2, 3, 1, 1) ) # 解码器 self.decoder nn.Sequential( nn.Conv2d(latent_channels, base_channels*4, 3, 1, 1), nn.GroupNorm(32, base_channels*4), nn.SiLU(), nn.ConvTranspose2d(base_channels*4, base_channels*2, 4, 2, 1), nn.GroupNorm(32, base_channels*2), nn.SiLU(), nn.ConvTranspose2d(base_channels*2, base_channels, 4, 2, 1), nn.GroupNorm(32, base_channels), nn.SiLU(), nn.ConvTranspose2d(base_channels, in_channels, 4, 2, 1), nn.Tanh() ) def encode(self, x): h self.encoder(x) mean, logvar torch.chunk(h, 2, dim1) return mean, logvar def decode(self, z): return self.decoder(z) def forward(self, x): mean, logvar self.encode(x) z self.reparameterize(mean, logvar) return self.decode(z), mean, logvar staticmethod def reparameterize(mean, logvar): std torch.exp(0.5 * logvar) eps torch.randn_like(std) return mean eps * std2.2 VAE训练技巧训练VAE时需要注意以下几点损失函数使用MSE重建损失和KL散度的组合def vae_loss(recon_x, x, mu, logvar): recon_loss F.mse_loss(recon_x, x, reductionsum) kl_loss -0.5 * torch.sum(1 logvar - mu.pow(2) - logvar.exp()) return recon_loss kl_loss * 0.0001学习率调度使用余弦退火学习率scheduler torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_maxepochs, eta_min1e-6)潜在空间缩放对潜在变量应用0.18215的缩放因子z self.reparameterize(mean, logvar) * 0.182153. 扩散模型核心实现3.1 噪声调度策略扩散模型需要定义噪声调度策略控制噪声如何随时间步增加def get_noise_schedule(schedule_type, timesteps, beta_start1e-4, beta_end2e-2): if schedule_type linear: return torch.linspace(beta_start, beta_end, timesteps) elif schedule_type cosine: steps timesteps 1 x torch.linspace(0, timesteps, steps) alphas_cumprod torch.cos((x / timesteps 0.008) / 1.008 * math.pi / 2) ** 2 alphas_cumprod alphas_cumprod / alphas_cumprod[0] betas 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1]) return torch.clip(betas, 0, 0.999) else: raise ValueError(fUnknown schedule type {schedule_type})3.2 UNet网络设计UNet是扩散模型的核心网络负责预测噪声class UNet(nn.Module): def __init__(self, in_channels3, out_channels3, base_channels64): super().__init__() # 下采样路径 self.down1 DownBlock(in_channels, base_channels) self.down2 DownBlock(base_channels, base_channels*2) self.down3 DownBlock(base_channels*2, base_channels*4) # 中间层 self.mid nn.Sequential( ResBlock(base_channels*4, base_channels*4), AttentionBlock(base_channels*4), ResBlock(base_channels*4, base_channels*4) ) # 上采样路径 self.up1 UpBlock(base_channels*8, base_channels*2) self.up2 UpBlock(base_channels*4, base_channels) self.up3 UpBlock(base_channels*2, base_channels) # 输出层 self.out nn.Conv2d(base_channels, out_channels, 3, 1, 1) def forward(self, x, t): # 下采样 h1 self.down1(x, t) h2 self.down2(h1, t) h3 self.down3(h2, t) # 中间层 h self.mid(h3) # 上采样 h self.up1(h, h2, t) h self.up2(h, h1, t) h self.up3(h, None, t) return self.out(h)4. 采样方法实现4.1 DDPM采样DDPM是最基础的采样方法def ddpm_sample(model, x, t, noise): alpha_t alpha[t].view(-1, 1, 1, 1) alpha_bar_t alpha_bar[t].view(-1, 1, 1, 1) eps model(x, t) x_prev (1 / torch.sqrt(alpha_t)) * ( x - ((1 - alpha_t) / torch.sqrt(1 - alpha_bar_t)) * eps ) torch.sqrt(beta[t]) * noise return x_prev4.2 DDIM采样DDIM通过非马尔可夫链加速采样def ddim_sample(model, x, t, t_prev, eta0): alpha_t alpha[t].view(-1, 1, 1, 1) alpha_bar_t alpha_bar[t].view(-1, 1, 1, 1) alpha_bar_prev alpha_bar[t_prev].view(-1, 1, 1, 1) eps model(x, t) x0 (x - torch.sqrt(1 - alpha_bar_t) * eps) / torch.sqrt(alpha_bar_t) sigma eta * torch.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar_t)) * torch.sqrt(1 - alpha_t) noise torch.randn_like(x) if t 0 else torch.zeros_like(x) x_prev torch.sqrt(alpha_bar_prev) * x0 \ torch.sqrt(1 - alpha_bar_prev - sigma**2) * eps \ sigma * noise return x_prev4.3 PLMS采样PLMS通过伪线性多步法进一步优化采样def plms_sample(model, x, t, t_prev, eps_history, eta0): alpha_t alpha[t].view(-1, 1, 1, 1) alpha_bar_t alpha_bar[t].view(-1, 1, 1, 1) alpha_bar_prev alpha_bar[t_prev].view(-1, 1, 1, 1) eps model(x, t) eps_history.append(eps) if len(eps_history) 1: eps_prime eps elif len(eps_history) 2: eps_prime (3 * eps - eps_history[-2]) / 2 elif len(eps_history) 3: eps_prime (23 * eps - 16 * eps_history[-2] 5 * eps_history[-3]) / 12 else: eps_prime (55 * eps - 59 * eps_history[-2] 37 * eps_history[-3] - 9 * eps_history[-4]) / 24 x0 (x - torch.sqrt(1 - alpha_bar_t) * eps_prime) / torch.sqrt(alpha_bar_t) sigma eta * torch.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar_t)) * torch.sqrt(1 - alpha_t) noise torch.randn_like(x) if t 0 else torch.zeros_like(x) x_prev torch.sqrt(alpha_bar_prev) * x0 \ torch.sqrt(1 - alpha_bar_prev - sigma**2) * eps_prime \ sigma * noise return x_prev5. 训练流程优化5.1 混合精度训练使用混合精度训练可以显著减少显存占用scaler torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): pred_noise model(noisy_images, timesteps) loss F.mse_loss(pred_noise, noise) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()5.2 EMA模型平滑使用EMA模型可以稳定训练过程class EMA: def __init__(self, beta0.9999): self.beta beta self.step 0 def update_model_average(self, ema_model, model): for ema_param, param in zip(ema_model.parameters(), model.parameters()): ema_param.data self.beta * ema_param.data (1 - self.beta) * param.data def step_ema(self, ema_model, model): if self.step 0: ema_model.load_state_dict(model.state_dict()) else: self.update_model_average(ema_model, model) self.step 15.3 分类器自由引导实现分类器自由引导以提升生成质量def forward(self, x, t, yNone, cfg_scale3.0): if y is None or cfg_scale 0: return self.model(x, t) # 无条件预测 uncond_out self.model(x, t, None) # 条件预测 cond_out self.model(x, t, y) # 线性插值 return uncond_out cfg_scale * (cond_out - uncond_out)6. 实际应用建议潜在空间压缩比选择1/8压缩比适合大多数场景高质量生成可使用1/4压缩比快速生成可使用1/16压缩比噪声调度策略对比策略类型训练稳定性生成质量适用场景Linear高一般快速实现Cosine中高高质量生成Sqrt低高研究实验采样方法选择指南DDPM基准方法速度慢但稳定DDIM20-50步即可获得不错结果PLMS10-30步达到最佳平衡显存优化技巧使用梯度检查点from torch.utils.checkpoint import checkpoint output checkpoint(self.model, x, t)减少批大小并使用梯度累积使用更小的UNet通道基数常见问题排查生成图像模糊检查VAE重建质量增加潜在空间维度训练不稳定降低学习率增加EMA衰减率模式崩溃增加分类器引导强度检查数据多样性在实际项目中我发现潜在扩散模型对超参数选择非常敏感。经过多次实验总结出以下经验当使用1/8压缩比时UNet的base_channels设置为64在质量和速度间取得较好平衡对于256x256图像生成余弦噪声调度配合DDIM采样在50步时效果最佳。另外分类器自由引导的scale参数设置在3-7之间通常能获得理想效果。