Skip to content

扩散模型 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 步。

相关资源

最近更新