07 - 连续时间扩散模型¶

import torch
import torch.nn as nn
import numpy as np
from scipy import integrate

class ContinuousDiffusion:
    """
    连续时间扩散模型
    """
    def __init__(self, sde_type='vp', beta_min=0.1, beta_max=20.0):
        """
        参数:
            sde_type: 'vp', 've', 或 'subvp'
            beta_min, beta_max: beta调度参数
        """
        self.sde_type = sde_type
        self.beta_min = beta_min
        self.beta_max = beta_max

    def beta(self, t):
        """beta(t)调度"""
        return self.beta_min + t * (self.beta_max - self.beta_min)

    def drift(self, x, t):
        """漂移系数 f(x, t)"""
        if self.sde_type == 'vp':
            return -0.5 * self.beta(t) * x
        elif self.sde_type == 've':
            return torch.zeros_like(x)
        elif self.sde_type == 'subvp':
            return -0.5 * self.beta(t) * x
        else:
            raise ValueError(f"Unknown SDE type: {self.sde_type}")

    def diffusion(self, t):
        """扩散系数 g(t)"""
        if self.sde_type == 'vp':
            return np.sqrt(self.beta(t))
        elif self.sde_type == 've':
            sigma = self.sigma(t)
            # g(t) = σ(t)√(2 ln(σ_max/σ_min))，由 d[σ²(t)]/dt 推导
            return sigma * np.sqrt(2 * np.log(self.beta_max / self.beta_min))
        elif self.sde_type == 'subvp':
            alpha = np.exp(-0.5 * self.beta_min * t - 0.25 * (self.beta_max - self.beta_min) * t**2)
            # g(t) = √(β(t)(1 - e^{-2∫β}))，其中 e^{-2∫β} = α⁴
            return np.sqrt(self.beta(t) * (1 - alpha**4))
        else:
            raise ValueError(f"Unknown SDE type: {self.sde_type}")

    def sigma(self, t):
        """VE-SDE的sigma(t)"""
        return self.beta_min * (self.beta_max / self.beta_min) ** t

    def marginal_prob(self, x0, t):
        """
        计算边缘分布 p(x(t) | x(0))

        返回:
            mean: 均值
            std: 标准差
        """
        if self.sde_type == 'vp':
            alpha = np.exp(-0.5 * self.beta_min * t - 0.25 * (self.beta_max - self.beta_min) * t**2)
            mean = alpha * x0
            std = np.sqrt(1 - alpha**2)
        elif self.sde_type == 've':
            mean = x0
            # 方差为 σ²(t) - σ²(0)，标准差取平方根
            std = np.sqrt(self.sigma(t)**2 - self.sigma(0)**2)
        elif self.sde_type == 'subvp':
            alpha = np.exp(-0.5 * self.beta_min * t - 0.25 * (self.beta_max - self.beta_min) * t**2)
            mean = alpha * x0
            # sub-VP-SDE: 方差 = (1-α²)²，标准差 = |1-α²| = 1-α² (因为α∈[0,1])
            # 注意：这里std是标准差，直接等于(1-α²)
            std = 1 - alpha**2
        else:
            raise ValueError(f"Unknown SDE type: {self.sde_type}")

        return mean, std

    def prior_sampling(self, shape):
        """从先验分布采样"""
        if self.sde_type == 'vp' or self.sde_type == 'subvp':
            return torch.randn(*shape)
        elif self.sde_type == 've':
            return torch.randn(*shape) * self.sigma(1.0)

    def forward_sde(self, x0, t, noise=None):
        """
        前向SDE: x(t) = mean + std * noise
        """
        mean, std = self.marginal_prob(x0, t)
        if noise is None:
            noise = torch.randn_like(x0)
        xt = mean + std * noise
        return xt, noise

    def reverse_sde(self, score_model, xt, t, dt, noise=None):
        """
        逆向SDE采样一步

        参数:
            score_model: 得分模型 s_θ(x, t)
            xt: 当前状态
            t: 当前时间
            dt: 时间步长（负值表示反向）
            noise: 可选的噪声
        """
        if noise is None:
            noise = torch.randn_like(xt)

        # 计算得分
        score = score_model(xt, t)

        # 漂移和扩散
        f = self.drift(xt, t)
        g = self.diffusion(t)

        # 逆向SDE: dx = (f - g² * score)dt + g * dw
        drift = f - g**2 * score
        diffusion = g

        # Euler-Maruyama
        x_prev = xt + drift * dt + diffusion * np.sqrt(-dt) * noise

        return x_prev

    def probability_flow_ode(self, score_model, xt, t, dt):
        """
        概率流ODE采样一步

        参数:
            score_model: 得分模型
            xt: 当前状态
            t: 当前时间
            dt: 时间步长（负值表示反向）
        """
        # 计算得分
        score = score_model(xt, t)

        # 漂移和扩散
        f = self.drift(xt, t)
        g = self.diffusion(t)

        # 概率流ODE: dx = (f - 0.5 * g² * score)dt
        dx = (f - 0.5 * g**2 * score) * dt

        x_prev = xt + dx

        return x_prev

    def sample_sde(self, score_model, shape, device='cuda', num_steps=1000, eps=1e-3):
        """
        使用SDE采样
        """
        score_model.eval()  # eval()评估模式

        # 从先验采样
        x = self.prior_sampling(shape).to(device)  # 移至GPU/CPU

        # 时间网格
        timesteps = torch.linspace(1.0, eps, num_steps)
        dt = -(1.0 - eps) / num_steps

        # 逆向采样
        with torch.no_grad():  # 禁用梯度计算，节省内存
            for i in range(num_steps):
                t = timesteps[i]
                t_batch = torch.ones(shape[0], device=device) * t
                x = self.reverse_sde(score_model, x, t_batch, dt)

        return x

    def sample_ode(self, score_model, shape, device='cuda', rtol=1e-5, atol=1e-5, method='RK45'):
        """
        使用ODE采样（自适应步长）
        """
        score_model.eval()

        # 初始条件
        x = self.prior_sampling(shape).to(device)

        # 定义ODE
        def ode_func(t, x):
            x = torch.tensor(x, device=device, dtype=torch.float32).reshape(shape)  # 重塑张量形状
            t = torch.ones(shape[0], device=device) * t

            with torch.no_grad():
                score = score_model(x, t)
                f = self.drift(x, t.item())  # 将单元素张量转为Python数值
                g = self.diffusion(t.item())
                dx = (f - 0.5 * g**2 * score).cpu().numpy().flatten()

            return dx

        # 求解ODE
        solution = integrate.solve_ivp(
            ode_func,
            [1.0, 1e-3],
            x.cpu().numpy().flatten(),
            rtol=rtol,
            atol=atol,
            method=method,
        )

        x_final = torch.tensor(solution.y[:, -1], device=device, dtype=torch.float32).reshape(shape)

        return x_final

# 得分模型示例
class ScoreNet(nn.Module):  # 继承nn.Module定义网络层
    """简化的得分网络"""
    def __init__(self, channels=3, time_emb_dim=256):
        super().__init__()  # super()调用父类方法
        self.time_embed = nn.Sequential(
            nn.Linear(1, time_emb_dim),
            nn.SiLU(),
            nn.Linear(time_emb_dim, time_emb_dim),
        )

        # 简化的UNet结构
        self.encoder = nn.Sequential(
            nn.Conv2d(channels, 64, 3, padding=1),
            nn.SiLU(),
            nn.Conv2d(64, 128, 3, padding=1),
            nn.SiLU(),
        )

        self.middle = nn.Sequential(
            nn.Conv2d(128, 128, 3, padding=1),
            nn.SiLU(),
        )

        self.decoder = nn.Sequential(
            nn.Conv2d(128, 64, 3, padding=1),
            nn.SiLU(),
            nn.Conv2d(64, channels, 3, padding=1),
        )

    def forward(self, x, t):
        # 时间嵌入
        if t.dim() == 0:
            t = t.unsqueeze(0).expand(x.shape[0])  # unsqueeze增加一个维度
        t_emb = self.time_embed(t.view(-1, 1).float())

        # 前向传播
        h = self.encoder(x)
        h = self.middle(h)
        h = self.decoder(h)

        return h

# 使用示例
if __name__ == "__main__":
    # 创建扩散模型
    diffusion = ContinuousDiffusion(sde_type='vp', beta_min=0.1, beta_max=20.0)

    # 创建得分模型
    score_model = ScoreNet(channels=3)

    # 测试前向过程
    x0 = torch.randn(4, 3, 32, 32)
    t = torch.tensor([0.5, 0.5, 0.5, 0.5])
    xt, noise = diffusion.forward_sde(x0, t)

    print(f"x0 shape: {x0.shape}")
    print(f"xt shape: {xt.shape}")
    print(f"noise shape: {noise.shape}")

    # 测试采样（需要训练好的模型）
    # samples = diffusion.sample_sde(score_model, shape=(4, 3, 32, 32))
    # print(f"Samples shape: {samples.shape}")

# 核心采样循环
def sample_sde(score_model, shape):
    x = prior_sampling(shape)
    for t in reversed(timesteps):
        score = score_model(x, t)
        x = reverse_sde_step(x, t, score)
    return x