03-扩散模型基础¶

import torch
import math

def linear_beta_schedule(T, beta_min=1e-4, beta_max=0.02):
    """线性噪声调度"""
    return torch.linspace(beta_min, beta_max, T)

def cosine_beta_schedule(T, s=0.008):
    """余弦噪声调度（Improved DDPM）"""
    steps = torch.arange(T + 1, dtype=torch.float64)
    alphas_cumprod = torch.cos(((steps / T) + s) / (1 + s) * math.pi * 0.5) ** 2
    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])  # 切片操作取子序列
    return torch.clip(betas, 0.0001, 0.9999).float()  # 链式调用，连续执行多个方法

def classifier_free_guidance(model, x_t, t, cond, w=7.5, uncond=None):
    """Classifier-Free Guidance采样"""
    # 有条件预测
    eps_cond = model(x_t, t, cond)
    # 无条件预测
    eps_uncond = model(x_t, t, uncond)
    # CFG组合
    eps_guided = (1 + w) * eps_cond - w * eps_uncond
    return eps_guided

class JointAttention(nn.Module):  # 继承nn.Module定义神经网络层
    """MMDiT的联合注意力机制"""
    def __init__(self, d_model, num_heads):  # __init__构造方法，创建对象时自动调用
        super().__init__()  # super()调用父类方法
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads
        self.to_q = nn.Linear(d_model, d_model)
        self.to_k = nn.Linear(d_model, d_model)
        self.to_v = nn.Linear(d_model, d_model)
        self.scale = self.head_dim ** -0.5

    def forward(self, img, txt):
        """
        img: (batch, img_tokens, d_model) - 图像tokens
        txt: (batch, txt_tokens, d_model) - 文本tokens
        """
        # 拼接图像和文本tokens
        x = torch.cat([img, txt], dim=1)  # (batch, img_tokens + txt_tokens, d_model)

        # 计算Q, K, V
        q = self.to_q(x)
        k = self.to_k(x)
        v = self.to_v(x)

        # 分割为多头
        q = q.view(q.size(0), -1, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(k.size(0), -1, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(v.size(0), -1, self.num_heads, self.head_dim).transpose(1, 2)

        # 自注意力（所有tokens互相attend）
        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        out = attn @ v

        # 拆分回图像和文本
        out = out.transpose(1, 2).contiguous().view(out.size(0), -1, self.d_model)
        img_out, txt_out = out.split([img.size(1), txt.size(1)], dim=1)

        return img_out, txt_out

# Flow Matching的核心思想（简化版）
class FlowMatchingLoss(nn.Module):
    """Flow Matching损失函数"""
    def forward(self, model, x_0, t):
        """
        x_0: 原始数据
        t: 时间步
        """
        # 1. 采样噪声和中间状态
        eps = torch.randn_like(x_0)
        x_t = (1 - t) * x_0 + t * eps  # 简单线性插值

        # 2. 预测速度场（velocity field）
        v_pred = model(x_t, t)

        # 3. 计算目标速度场
        v_target = eps - x_0  # 从x_0到噪声的方向

        # 4. MSE损失
        loss = F.mse_loss(v_pred, v_target)
        return loss

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import math

# ========================================
# 1. 噪声调度器
# ========================================
class NoiseScheduler:
    """管理DDPM的噪声调度参数"""

    def __init__(self, T=1000, beta_min=1e-4, beta_max=0.02, device='cpu'):
        self.T = T
        self.device = device

        # 线性beta调度
        self.betas = torch.linspace(beta_min, beta_max, T).to(device)
        self.alphas = 1.0 - self.betas
        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
        self.alphas_cumprod_prev = F.pad(self.alphas_cumprod[:-1], (1, 0), value=1.0)

        # 预计算常用系数
        self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod)
        self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - self.alphas_cumprod)
        self.sqrt_recip_alphas = torch.sqrt(1.0 / self.alphas)

        # 逆向过程的后验方差
        self.posterior_variance = (
            self.betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
        )

    def add_noise(self, x_0, t, noise=None):
        """前向过程：给x_0添加噪声得到x_t"""
        if noise is None:
            noise = torch.randn_like(x_0)

        sqrt_alpha = self.sqrt_alphas_cumprod[t].view(-1, 1, 1, 1)  # view重塑张量形状（要求内存连续）
        sqrt_one_minus_alpha = self.sqrt_one_minus_alphas_cumprod[t].view(-1, 1, 1, 1)

        x_t = sqrt_alpha * x_0 + sqrt_one_minus_alpha * noise
        return x_t, noise

    @torch.no_grad()  # 禁用梯度计算，节省内存
    def sample_step(self, model, x_t, t):
        """逆向过程：单步去噪 x_t -> x_{t-1}"""
        t_tensor = torch.full((x_t.shape[0],), t, device=self.device, dtype=torch.long)

        # 预测噪声
        predicted_noise = model(x_t, t_tensor)

        # 计算均值
        beta_t = self.betas[t]
        sqrt_recip_alpha_t = self.sqrt_recip_alphas[t]
        sqrt_one_minus_alpha_cumprod_t = self.sqrt_one_minus_alphas_cumprod[t]

        mean = sqrt_recip_alpha_t * (
            x_t - beta_t / sqrt_one_minus_alpha_cumprod_t * predicted_noise
        )

        if t > 0:
            noise = torch.randn_like(x_t)
            variance = torch.sqrt(self.posterior_variance[t])
            return mean + variance * noise
        else:
            return mean

# ========================================
# 2. 时间步嵌入
# ========================================
class SinusoidalPositionEmbedding(nn.Module):
    """正弦位置编码，将时间步t编码为向量"""

    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, t):
        device = t.device
        half_dim = self.dim // 2
        emb = math.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
        emb = t[:, None].float() * emb[None, :]
        emb = torch.cat([emb.sin(), emb.cos()], dim=-1)
        return emb

# ========================================
# 3. 简单UNet去噪网络
# ========================================
class SimpleUNet(nn.Module):
    """简化版UNet，用于MNIST去噪"""

    def __init__(self, in_channels=1, time_dim=256):
        super().__init__()

        # 时间步嵌入
        self.time_mlp = nn.Sequential(
            SinusoidalPositionEmbedding(time_dim),
            nn.Linear(time_dim, time_dim),
            nn.GELU(),
            nn.Linear(time_dim, time_dim),
        )

        # 编码器（下采样）
        self.enc1 = self._make_block(in_channels, 64)
        self.enc2 = self._make_block(64, 128)
        self.pool = nn.MaxPool2d(2)

        # 瓶颈层
        self.bottleneck = self._make_block(128, 256)
        self.time_proj = nn.Linear(time_dim, 256)

        # 解码器（上采样）
        self.up2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
        self.dec2 = self._make_block(256, 128)  # 256 = 128(up) + 128(skip)
        self.up1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
        self.dec1 = self._make_block(128, 64)   # 128 = 64(up) + 64(skip)

        # 输出层
        self.out = nn.Conv2d(64, in_channels, kernel_size=1)

    def _make_block(self, in_ch, out_ch):
        return nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, padding=1),
            nn.GroupNorm(8, out_ch),
            nn.GELU(),
            nn.Conv2d(out_ch, out_ch, 3, padding=1),
            nn.GroupNorm(8, out_ch),
            nn.GELU(),
        )

    def forward(self, x, t):
        # 时间步嵌入
        t_emb = self.time_mlp(t)  # (B, time_dim)

        # 编码器
        e1 = self.enc1(x)                    # (B, 64, 28, 28)
        e2 = self.enc2(self.pool(e1))        # (B, 128, 14, 14)

        # 瓶颈层 + 时间条件注入
        b = self.bottleneck(self.pool(e2))   # (B, 256, 7, 7)
        t_proj = self.time_proj(t_emb)[:, :, None, None]
        b = b + t_proj                        # 时间条件加入

        # 解码器 + 跳跃连接
        d2 = self.up2(b)                      # (B, 128, 14, 14)
        d2 = self.dec2(torch.cat([d2, e2], dim=1))
        d1 = self.up1(d2)                     # (B, 64, 28, 28)
        d1 = self.dec1(torch.cat([d1, e1], dim=1))

        return self.out(d1)

# ========================================
# 4. 训练循环
# ========================================
def train_ddpm(
    epochs=10,
    batch_size=128,
    lr=2e-4,
    T=1000,
    device='cuda' if torch.cuda.is_available() else 'cpu',
):
    """训练DDPM模型"""
    print(f"Using device: {device}")

    # 数据集
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize([0.5], [0.5]),  # 归一化到 [-1, 1]
    ])
    dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=2)  # DataLoader批量加载数据，支持shuffle和多进程

    # 模型和优化器
    model = SimpleUNet(in_channels=1).to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    scheduler = NoiseScheduler(T=T, device=device)

    print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")

    # 训练
    for epoch in range(epochs):
        total_loss = 0
        for batch_idx, (images, _) in enumerate(dataloader):  # enumerate同时获取索引和元素
            images = images.to(device)

            # 随机采样时间步
            t = torch.randint(0, T, (images.shape[0],), device=device)

            # 前向加噪
            x_t, noise = scheduler.add_noise(images, t)

            # 预测噪声
            predicted_noise = model(x_t, t)

            # MSE损失
            loss = F.mse_loss(predicted_noise, noise)

            # 反向传播
            optimizer.zero_grad()
            loss.backward()  # 反向传播计算梯度
            # 梯度裁剪（稳定训练）
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()  # 根据梯度更新模型参数

            total_loss += loss.item()  # .item()将单元素张量转为Python数值

        avg_loss = total_loss / len(dataloader)
        print(f"Epoch [{epoch+1}/{epochs}], Loss: {avg_loss:.4f}")

    return model, scheduler

# ========================================
# 5. 采样/生成
# ========================================
@torch.no_grad()
def generate_samples(
    model, scheduler, n_samples=16, device='cuda',
    image_size=28, channels=1
):
    """从训练好的DDPM生成图像"""
    model.eval()  # eval()开启评估模式（关闭Dropout等）

    # 从纯噪声开始
    x = torch.randn(n_samples, channels, image_size, image_size).to(device)

    # 逐步去噪 T -> 0
    for t in reversed(range(scheduler.T)):
        x = scheduler.sample_step(model, x, t)

    # 将 [-1, 1] 映射回 [0, 1]
    x = (x + 1) / 2
    x = x.clamp(0, 1)

    return x

# ========================================
# 主入口
# ========================================
if __name__ == '__main__':
    # 训练模型
    model, scheduler = train_ddpm(epochs=10, batch_size=128)

    # 生成样本
    device = next(model.parameters()).device
    samples = generate_samples(model, scheduler, n_samples=16, device=device)

    print(f"Generated {samples.shape[0]} samples, shape: {samples.shape}")

    # 可选：可视化（需要matplotlib）
    # import matplotlib.pyplot as plt
    # fig, axes = plt.subplots(4, 4, figsize=(8, 8))
    # for i, ax in enumerate(axes.flat):
    #     ax.imshow(samples[i, 0].cpu().numpy(), cmap='gray')
    #     ax.axis('off')
    # plt.savefig('ddpm_samples.png')
    # plt.show()