1
算法：DDPM训练
2
输入：数据集 {x_0}, 噪声调度 {β_t}, 总步数 T
3
输出：去噪网络 ε_θ
4

5
注：时间步索引 t ∈ {0, 1, ..., T-1}
6

7
重复：
8
    1. 从数据集采样 x_0 ~ q(x_0)
9
    2. 随机采样时间步 t ~ Uniform({0, 1, ..., T-1})
10
    3. 采样噪声 ε ~ N(0, I)
11
    4. 计算 x_t = √(ᾱ_t) x_0 + √(1-ᾱ_t) ε
12
    5. 计算损失 L = ||ε - ε_θ(x_t, t)||²
13
    6. 梯度下降更新 θ
14
直到收敛

1
输入: x_t (噪声图像) + t (时间步)
2

3
编码器:
4
  Conv Block 1 (64×64) + Time Embed → 特征1
5
  ↓ 下采样
6
  Conv Block 2 (32×32) + Time Embed → 特征2
7
  ↓ 下采样
8
  Conv Block 3 (16×16) + Time Embed → 特征3
9
  ↓ 下采样
10

11
瓶颈层:
12
  Conv Block 4 (8×8) + Self-Attention + Time Embed
13

14
解码器:
15
  ↑ 上采样 + 拼接特征3
16
  Conv Block 5 (16×16) + Time Embed
17
  ↑ 上采样 + 拼接特征2
18
  Conv Block 6 (32×32) + Time Embed
19
  ↑ 上采样 + 拼接特征1
20
  Conv Block 7 (64×64) + Time Embed
21

22
输出: ε_pred (预测的噪声)

1
import torch
2
import torch.nn as nn
3
import torch.nn.functional as F
4

5
class SimpleDiffusion:
6
    def __init__(self, num_steps=1000, beta_start=1e-4, beta_end=2e-2, device="cuda"):
7
        self.num_steps = int(num_steps)
8
        self.device = device
9

10
        # 线性噪声调度（教学默认）
11
        betas = torch.linspace(beta_start, beta_end, self.num_steps, device=device)
12
        alphas = 1.0 - betas
13
        alpha_bars = torch.cumprod(alphas, dim=0)
14

15
        # 保存为成员（在真实工程里可注册为 buffer）
16
        self.betas = betas
17
        self.alphas = alphas
18
        self.alpha_bars = alpha_bars
19

20
    def _extract(self, a, t, x_shape):
21
        # 从长度为 T 的向量 a 中按 batch 的 t 抽取对应元素，并 reshape 成可广播形状
22
        b = t.shape[0]
23
        out = a.gather(0, t).view(b, *([1] * (len(x_shape) - 1)))
24
        return out
25

26
    def forward_diffusion(self, x0, t):
27
        # 前向扩散：从 x0 直接得到 xt（一步到位）
28
        # 注意：x0 建议预处理到 [-1, 1]
29
        noise = torch.randn_like(x0)
30
        alpha_bar_t = self._extract(self.alpha_bars, t, x0.shape)
31
        xt = torch.sqrt(alpha_bar_t) * x0 + torch.sqrt(1.0 - alpha_bar_t) * noise
32
        return xt, noise
33

34
    @torch.no_grad()
35
    def reverse_step(self, model, xt, t):
36
        # 反向一步：xt -> x_{t-1}
37
        # ε-prediction 参数化 + 后验方差 beta_tilde
38
        eps_pred = model(xt, t)
39

40
        alpha_t = self._extract(self.alphas, t, xt.shape)
41
        alpha_bar_t = self._extract(self.alpha_bars, t, xt.shape)
42
        beta_t = self._extract(self.betas, t, xt.shape)
43

44
        # DDPM posterior mean (ε-parameterization)
45
        mean = (1.0 / torch.sqrt(alpha_t)) * (xt - (beta_t / torch.sqrt(1.0 - alpha_bar_t)) * eps_pred)
46

47
        # posterior variance beta_tilde
48
        t_prev = (t - 1).clamp(min=0)
49
        alpha_bar_prev = self._extract(self.alpha_bars, t_prev, xt.shape)
50
        beta_tilde = ((1.0 - alpha_bar_prev) / (1.0 - alpha_bar_t)) * beta_t
51
        sigma = torch.sqrt(beta_tilde)
52

53
        # 当 t==0 时不再加噪声
54
        noise = torch.randn_like(xt)
55
        mask = (t > 0).float().view(t.shape[0], *([1] * (len(xt.shape) - 1)))
56
        return mean + mask * sigma * noise
57

58
def train_step(model, diffusion: SimpleDiffusion, x0, optimizer):
59
    # 注意：x0 建议预处理到 [-1, 1]
60
    b = x0.shape[0]
61
    device = x0.device
62
    t = torch.randint(0, diffusion.num_steps, (b,), device=device)
63

64
    xt, noise = diffusion.forward_diffusion(x0, t)
65
    noise_pred = model(xt, t)
66

67
    loss = F.mse_loss(noise_pred, noise)
68
    optimizer.zero_grad()
69
    loss.backward()
70
    optimizer.step()
71
    return float(loss.item())
72

73
@torch.no_grad()
74
def sample(model, diffusion: SimpleDiffusion, shape, device):
75
    # 从纯噪声生成图像，返回值在 [-1, 1]，后处理：x = (x + 1) * 127.5
76
    x = torch.randn(shape, device=device)
77
    for t in reversed(range(diffusion.num_steps)):
78
        t_batch = torch.full((shape[0],), t, dtype=torch.long, device=device)
79
        x = diffusion.reverse_step(model, x, t_batch)
80
    return x

1
import torch
2
import torch.nn as nn
3
import torch.nn.functional as F
4

5
class SimpleUNet(nn.Module):
6
    # 教学版 U-Net：预测噪声 ε
7
    # 显式注入 time embedding，并使用 GroupNorm + SiLU（更贴近真实扩散 U-Net）
8
    def __init__(self, in_channels=3, base_channels=64, time_dim=256):
9
        super().__init__()
10
        self.time_dim = time_dim
11

12
        # time embedding: sinusoidal -> MLP
13
        self.time_mlp = nn.Sequential(
14
            nn.Linear(time_dim, time_dim * 4),
15
            nn.SiLU(),
16
            nn.Linear(time_dim * 4, time_dim),
17
        )
18

19
        # encoder
20
        self.conv1 = nn.Conv2d(in_channels, base_channels, 3, padding=1)
21
        self.gn1 = nn.GroupNorm(8, base_channels)
22
        self.tproj1 = nn.Linear(time_dim, base_channels)
23

24
        self.conv2 = nn.Conv2d(base_channels, base_channels * 2, 3, padding=1)
25
        self.gn2 = nn.GroupNorm(8, base_channels * 2)
26
        self.tproj2 = nn.Linear(time_dim, base_channels * 2)
27

28
        self.conv3 = nn.Conv2d(base_channels * 2, base_channels * 4, 3, padding=1)
29
        self.gn3 = nn.GroupNorm(8, base_channels * 4)
30
        self.tproj3 = nn.Linear(time_dim, base_channels * 4)
31

32
        # bottleneck
33
        self.conv4 = nn.Conv2d(base_channels * 4, base_channels * 8, 3, padding=1)
34
        self.gn4 = nn.GroupNorm(8, base_channels * 8)
35
        self.tproj4 = nn.Linear(time_dim, base_channels * 8)
36

37
        # decoder
38
        self.dconv3 = nn.Conv2d(base_channels * 8 + base_channels * 4, base_channels * 4, 3, padding=1)
39
        self.dgn3 = nn.GroupNorm(8, base_channels * 4)
40
        self.dtproj3 = nn.Linear(time_dim, base_channels * 4)
41

42
        self.dconv2 = nn.Conv2d(base_channels * 4 + base_channels * 2, base_channels * 2, 3, padding=1)
43
        self.dgn2 = nn.GroupNorm(8, base_channels * 2)
44
        self.dtproj2 = nn.Linear(time_dim, base_channels * 2)
45

46
        self.dconv1 = nn.Conv2d(base_channels * 2 + base_channels, base_channels, 3, padding=1)
47
        self.dgn1 = nn.GroupNorm(8, base_channels)
48
        self.dtproj1 = nn.Linear(time_dim, base_channels)
49

50
        self.out = nn.Conv2d(base_channels, in_channels, 1)
51

52
    def forward(self, x, t):
53
        # t: (B,)
54
        temb = self.get_time_embedding(t, self.time_dim)
55
        temb = self.time_mlp(temb)
56

57
        def add_time(h, proj):
58
            return h + proj(temb).view(h.shape[0], -1, 1, 1)
59

60
        # encoder
61
        e1 = F.silu(add_time(self.gn1(self.conv1(x)), self.tproj1))
62
        x2 = F.avg_pool2d(e1, 2)
63

64
        e2 = F.silu(add_time(self.gn2(self.conv2(x2)), self.tproj2))
65
        x3 = F.avg_pool2d(e2, 2)
66

67
        e3 = F.silu(add_time(self.gn3(self.conv3(x3)), self.tproj3))
68
        x4 = F.avg_pool2d(e3, 2)
69

70
        # bottleneck
71
        b = F.silu(add_time(self.gn4(self.conv4(x4)), self.tproj4))
72

73
        # decoder
74
        d3 = F.interpolate(b, scale_factor=2, mode="nearest")
75
        d3 = torch.cat([d3, e3], dim=1)
76
        d3 = F.silu(add_time(self.dgn3(self.dconv3(d3)), self.dtproj3))
77

78
        d2 = F.interpolate(d3, scale_factor=2, mode="nearest")
79
        d2 = torch.cat([d2, e2], dim=1)
80
        d2 = F.silu(add_time(self.dgn2(self.dconv2(d2)), self.dtproj2))
81

82
        d1 = F.interpolate(d2, scale_factor=2, mode="nearest")
83
        d1 = torch.cat([d1, e1], dim=1)
84
        d1 = F.silu(add_time(self.dgn1(self.dconv1(d1)), self.dtproj1))
85

86
        return self.out(d1)
87

88
    @staticmethod
89
    def get_time_embedding(t, dim=256):
90
        # sinusoidal time embedding: (B, dim)
91
        half = dim // 2
92
        freqs = torch.exp(
93
            -torch.log(torch.tensor(10000.0, device=t.device)) * torch.arange(half, device=t.device) / (half - 1)
94
        )
95
        args = t.float()[:, None] * freqs[None, :]
96
        emb = torch.cat([torch.sin(args), torch.cos(args)], dim=-1)
97
        return emb

1
文本 → CLIP编码器 → 文本嵌入
2
图像 → VAE编码器 → 潜在表示 → U-Net扩散 → 潜在表示 → VAE解码器 → 图像