2D扩散模型入门指南：从数学基础到PyTorch实战

1次阅读

没有评论

共计 1792 个字符，预计需要花费 5 分钟才能阅读完成。

在图像生成领域，我们熟悉的 GAN 和 VAE 各有优缺点：

GAN：生成质量高但训练不稳定，容易模式坍塌
VAE：训练稳定但生成图像往往模糊

扩散模型则提供了一种新思路——通过逐步加噪和去噪的过程，既保持了生成质量，又让训练过程更加可控。就像画家先打草稿再逐步细化，扩散模型通过多个时间步的迭代达到理想效果。

扩散模型的核心是 前向过程 和反向过程：

前向过程（加噪）：
定义为一个马尔可夫链，逐步向数据添加高斯噪声：
$$q(\mathbf{x}t|\mathbf{x}}) = \mathcal{N}(\mathbf{xt; \sqrt{1-\beta_t}\mathbf{x})$$
其中 $\beta_t$ 是噪声调度参数}, \beta_t\mathbf{I
噪声调度策略：
通常采用线性或余弦调度，控制噪声增加的节奏。余弦调度效果更好：
$$\beta_t = \text{clip}(1-\frac{\alpha_t}{\alpha_{t-1}}, 0.999)$$
$$\alpha_t = \frac{\cos(t/T+s}{1+s}\pi/2)^2$$

import torch
def gaussian_noise(x, t, betas):
    """为输入 x 在时间步 t 添加噪声"""
    sqrt_alpha = torch.sqrt(1 - betas[t])
    noise = torch.randn_like(x)
    return sqrt_alpha * x + (1 - sqrt_alpha) * noise

class TimeEmbedding(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim
        # 使用正弦位置编码
        half_dim = dim // 2
        emb = math.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim) * -emb)
        self.register_buffer('emb', emb)

    def forward(self, t):
        emb = t[:, None] * self.emb[None, :]
        return torch.cat([emb.sin(), emb.cos()], dim=-1)

class UNet(nn.Module):
    def __init__(self):
        super().__init__()
        # 下采样路径
        self.down1 = nn.Sequential(nn.Conv2d(3, 64, 3, padding=1),
            nn.GroupNorm(8, 64),
            nn.SiLU())
        # 上采样路径（包含 skip connection）self.up1 = nn.Sequential(nn.Conv2d(128, 64, 3, padding=1),
            nn.GroupNorm(8, 64),
            nn.SiLU())
        # ... 其他层定义

    def forward(self, x, t):
        # 时间嵌入
        t_emb = self.time_mlp(t)
        # 下采样
        h1 = self.down1(x)
        # ... 中间层
        # 上采样 +skip connection
        h = self.up1(torch.cat([h, h_skip], dim=1))
        return h