扩散模型 PyTorch 实现
扩散模型通过逐步去噪生成高质量图像,已成为图像生成的核心架构。
DDPM 前向与反向过程
python
class GaussianDiffusion:
def __init__(self, timesteps=1000, beta_start=1e-4, beta_end=0.02):
self.betas = torch.linspace(beta_start, beta_end, timesteps)
self.alphas = 1.0 - self.betas
self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
def q_sample(self, x_0, t, noise=None):
"""前向加噪过程"""
if noise is None:
noise = torch.randn_like(x_0)
sqrt_alpha = self.alphas_cumprod[t].sqrt().view(-1, 1, 1, 1)
sqrt_one_minus = (1 - self.alphas_cumprod[t]).sqrt().view(-1, 1, 1, 1)
return sqrt_alpha * x_0 + sqrt_one_minus * noise
def p_sample(self, model, x_t, t):
"""反向去噪一步"""
pred_noise = model(x_t, t)
alpha = self.alphas[t].view(-1, 1, 1, 1)
alpha_cumprod = self.alphas_cumprod[t].view(-1, 1, 1, 1)
pred_x0 = (x_t - (1 - alpha_cumprod).sqrt() * pred_noise) / alpha_cumprod.sqrt()
return pred_x0扩散模型的优势
相比 GAN,扩散模型训练更稳定、生成多样性更好。缺点是推理速度慢——需要数百步去噪。DDIM 和 LCM 等方法将步数压缩到 4-8 步。