05 - 采样算法详解¶

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt

def get_schedule(T, beta_start=0.0001, beta_end=0.02):
    """
    创建噪声调度表

    参数:
        T: 总步数
        beta_start: 初始beta值
        beta_end: 最终beta值

    返回:
        alphas, betas, alphas_cumprod
    """
    betas = torch.linspace(beta_start, beta_end, T)
    alphas = 1 - betas
    alphas_cumprod = torch.cumprod(alphas, dim=0)

    return alphas, betas, alphas_cumprod

def ddpm_sample(model, x_T, T, alphas, betas, alphas_cumprod, device='cpu'):
    """
    DDPM采样算法

    参数:
        model: 训练好的去噪模型
        x_T: 初始噪声 [batch_size, channels, height, width]
        T: 总步数
        alphas, betas, alphas_cumprod: 调度表
        device: 设备

    返回:
        生成的图像 x_0
    """
    model.eval()  # eval()评估模式
    x_t = x_T.to(device)  # 移至GPU/CPU

    # 预计算一些值
    alphas_cumprod_prev = torch.cat([torch.tensor([1.0]), alphas_cumprod[:-1]])  # torch.cat沿已有维度拼接张量

    with torch.no_grad():  # 禁用梯度计算，节省内存
        for t in reversed(range(T)):
            # 当前时间步的参数
            alpha_t = alphas[t]
            beta_t = betas[t]
            alpha_t_bar = alphas_cumprod[t]
            alpha_t_bar_prev = alphas_cumprod_prev[t]

            # 预测噪声
            t_tensor = torch.full((x_t.shape[0],), t, device=device, dtype=torch.long)
            predicted_noise = model(x_t, t_tensor)

            # 计算均值
            sqrt_recip_alpha_t = 1 / torch.sqrt(alpha_t)
            sqrt_alpha_t_bar = torch.sqrt(alpha_t_bar)
            sqrt_one_minus_alpha_t_bar = torch.sqrt(1 - alpha_t_bar)

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

            # 添加随机噪声（最后一步不添加）
            if t > 0:
                posterior_variance = beta_t * (1 - alpha_t_bar_prev) / (1 - alpha_t_bar)
                noise = torch.randn_like(x_t)
                x_t = mean + torch.sqrt(posterior_variance) * noise
            else:
                x_t = mean

    return x_t

# 示例：简单的UNet模型
class SimpleUNet(nn.Module):  # 继承nn.Module定义网络层
    def __init__(self, in_channels=3, out_channels=3):
        super().__init__()  # super()调用父类方法
        self.conv1 = nn.Conv2d(in_channels, 64, 3, padding=1)
        self.conv2 = nn.Conv2d(64, 128, 3, padding=1)
        self.conv3 = nn.Conv2d(128, 64, 3, padding=1)
        self.conv4 = nn.Conv2d(64, out_channels, 3, padding=1)
        self.relu = nn.ReLU()

    def forward(self, x, t):
        # 简化版，实际UNet更复杂
        x = self.relu(self.conv1(x))
        x = self.relu(self.conv2(x))
        x = self.relu(self.conv3(x))
        x = self.conv4(x)
        return x

# 测试采样
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"使用设备: {device}")

# 创建模型和调度表
T = 1000
model = SimpleUNet().to(device)
alphas, betas, alphas_cumprod = get_schedule(T)

# 生成初始噪声
batch_size = 4
x_T = torch.randn(batch_size, 3, 32, 32)

# 采样
print(f"开始采样，总步数: {T}")
x_0 = ddpm_sample(model, x_T, T, alphas, betas, alphas_cumprod, device)
print(f"采样完成，输出形状: {x_0.shape}")

def sample_with_steps(model, x_T, steps, alphas, betas, alphas_cumprod, device='cpu'):
    """
    使用指定步数进行采样

    参数:
        steps: 采样步数列表，例如 [10, 50, 100, 500, 1000]

    返回:
        不同步数下的生成结果
    """
    results = {}
    T = len(alphas)

    for num_steps in steps:
        # 计算子采样间隔
        step_indices = torch.linspace(0, T-1, num_steps, dtype=torch.long)

        x_t = x_T.clone()
        alphas_cumprod_prev = torch.cat([torch.tensor([1.0]), alphas_cumprod[:-1]])

        with torch.no_grad():
            for i, t in enumerate(reversed(step_indices)):  # enumerate同时获取索引和元素
                t = t.item()  # 将单元素张量转为Python数值
                alpha_t = alphas[t]
                beta_t = betas[t]
                alpha_t_bar = alphas_cumprod[t]
                alpha_t_bar_prev = alphas_cumprod_prev[t]

                # 预测噪声
                t_tensor = torch.full((x_t.shape[0],), t, device=device, dtype=torch.long)
                predicted_noise = model(x_t, t_tensor)

                # 计算均值
                sqrt_recip_alpha_t = 1 / torch.sqrt(alpha_t)
                sqrt_one_minus_alpha_t_bar = torch.sqrt(1 - alpha_t_bar)

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

                # 添加噪声
                if i < len(step_indices) - 1:  # 不是最后一步
                    posterior_variance = beta_t * (1 - alpha_t_bar_prev) / (1 - alpha_t_bar)
                    noise = torch.randn_like(x_t)
                    x_t = mean + torch.sqrt(posterior_variance) * noise
                else:
                    x_t = mean

        results[num_steps] = x_t.detach().cpu()  # 分离计算图，不参与梯度计算

    return results

# 测试不同步数
steps_list = [10, 50, 100, 500, 1000]
results = sample_with_steps(model, x_T, steps_list, alphas, betas, alphas_cumprod, device)

# 可视化
fig, axes = plt.subplots(1, len(steps_list), figsize=(15, 3))

for idx, steps in enumerate(steps_list):
    img = results[steps][0]  # 取第一张图
    img = img.permute(1, 2, 0).numpy()
    img = (img - img.min()) / (img.max() - img.min())  # 归一化

    axes[idx].imshow(img)
    axes[idx].set_title(f'{steps} 步')
    axes[idx].axis('off')

plt.tight_layout()
plt.show()

def ddim_sample(model, x_T, T, alphas, alphas_cumprod, num_steps=50, device='cpu'):
    """
    DDIM采样算法

    参数:
        model: 训练好的去噪模型
        x_T: 初始噪声
        T: 原始总步数
        alphas, alphas_cumprod: 调度表
        num_steps: 实际采样步数
        device: 设备

    返回:
        生成的图像 x_0
    """
    model.eval()
    x_t = x_T.to(device)

    # 选择采样时间步
    step_indices = torch.linspace(0, T-1, num_steps, dtype=torch.long)

    with torch.no_grad():
        for i, t in enumerate(reversed(step_indices)):
            t = t.item()
            alpha_t = alphas[t]
            alpha_t_bar = alphas_cumprod[t]

            # 预测噪声
            t_tensor = torch.full((x_t.shape[0],), t, device=device, dtype=torch.long)
            predicted_noise = model(x_t, t_tensor)

            # 计算x_0的预测
            sqrt_alpha_t_bar = torch.sqrt(alpha_t_bar)
            sqrt_one_minus_alpha_t_bar = torch.sqrt(1 - alpha_t_bar)
            x_0_pred = (x_t - sqrt_one_minus_alpha_t_bar * predicted_noise) / sqrt_alpha_t_bar

            # 计算x_{t-1}
            if i < len(step_indices) - 1:  # 不是最后一步
                t_prev = step_indices[len(step_indices) - 2 - i].item()
                alpha_t_prev_bar = alphas_cumprod[t_prev]

                sqrt_alpha_t_prev_bar = torch.sqrt(alpha_t_prev_bar)
                sqrt_one_minus_alpha_t_prev_bar = torch.sqrt(1 - alpha_t_prev_bar)

                # DDIM更新
                x_t = sqrt_alpha_t_prev_bar * x_0_pred + sqrt_one_minus_alpha_t_prev_bar * predicted_noise
            else:
                x_t = x_0_pred

    return x_t

# 测试DDIM
print("测试DDIM采样...")
for num_steps in [10, 20, 50, 100]:
    x_0_ddim = ddim_sample(model, x_T, T, alphas, alphas_cumprod, num_steps, device)
    print(f"DDIM {num_steps} 步采样完成，输出形状: {x_0_ddim.shape}")

def get_cosine_schedule(T, s=0.008):
    """
    余弦噪声调度（Improved DDPM）

    参数:
        T: 总步数
        s: 偏移参数

    返回:
        alphas, betas, alphas_cumprod
    """
    steps = T + 1
    x = torch.linspace(0, T, steps)
    alphas_cumprod = torch.cos(((x / T) + s) / (1 + s) * torch.pi * 0.5) ** 2
    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]

    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
    alphas = 1 - betas

    # 裁剪beta值
    betas = torch.clip(betas, 0.0001, 0.9999)

    return alphas, betas, alphas_cumprod[:-1]

# 对比不同调度
T = 1000

# 线性调度
alphas_linear, betas_linear, alphas_cumprod_linear = get_schedule(T)

# 余弦调度
alphas_cosine, betas_cosine, alphas_cumprod_cosine = get_cosine_schedule(T)

# 可视化
plt.figure(figsize=(12, 4))

plt.subplot(1, 3, 1)
plt.plot(betas_linear.numpy(), label='线性')
plt.plot(betas_cosine.numpy(), label='余弦')
plt.xlabel('时间步 t')
plt.ylabel('β_t')
plt.title('Beta 调度')
plt.legend()
plt.grid(True, alpha=0.3)

plt.subplot(1, 3, 2)
plt.plot(alphas_cumprod_linear.numpy(), label='线性')
plt.plot(alphas_cumprod_cosine.numpy(), label='余弦')
plt.xlabel('时间步 t')
plt.ylabel('ᾱ_t')
plt.title('Alpha Bar 调度')
plt.legend()
plt.grid(True, alpha=0.3)

plt.subplot(1, 3, 3)
plt.plot(np.sqrt(1 - alphas_cumprod_linear.numpy()), label='线性')
plt.plot(np.sqrt(1 - alphas_cumprod_cosine.numpy()), label='余弦')
plt.xlabel('时间步 t')
plt.ylabel('√(1-ᾱ_t)')
plt.title('噪声水平')
plt.legend()
plt.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

def classifier_guided_sample(model, classifier, x_T, T, alphas, betas, alphas_cumprod,
                             guidance_scale=1.0, class_label=None, device='cpu'):
    """
    分类器引导采样

    参数:
        model: 去噪模型
        classifier: 分类器
        x_T: 初始噪声
        T: 总步数
        alphas, betas, alphas_cumprod: 调度表
        guidance_scale: 引导强度
        class_label: 目标类别
        device: 设备

    返回:
        生成的图像
    """
    model.eval()
    classifier.eval()
    x_t = x_T.to(device)

    alphas_cumprod_prev = torch.cat([torch.tensor([1.0]), alphas_cumprod[:-1]])

    with torch.no_grad():
        for t in reversed(range(T)):
            alpha_t = alphas[t]
            beta_t = betas[t]
            alpha_t_bar = alphas_cumprod[t]

            t_tensor = torch.full((x_t.shape[0],), t, device=device, dtype=torch.long)

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

            # 如果有分类器，计算梯度引导
            if classifier is not None and class_label is not None:
                # 计算分类器梯度
                x_t.requires_grad = True
                logits = classifier(x_t, t_tensor)
                log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
                target_log_prob = log_probs[:, class_label].sum()
                target_log_prob.backward()  # 反向传播计算梯度

                grad = x_t.grad
                x_t.requires_grad = False

                # 添加引导
                predicted_noise = predicted_noise - guidance_scale * grad * torch.sqrt(1 - alpha_t_bar)

            # 更新
            sqrt_recip_alpha_t = 1 / torch.sqrt(alpha_t)
            sqrt_one_minus_alpha_t_bar = torch.sqrt(1 - alpha_t_bar)

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

            if t > 0:
                alpha_t_bar_prev = alphas_cumprod_prev[t]
                posterior_variance = beta_t * (1 - alpha_t_bar_prev) / (1 - alpha_t_bar)
                noise = torch.randn_like(x_t)
                x_t = mean + torch.sqrt(posterior_variance) * noise
            else:
                x_t = mean

    return x_t

def classifier_free_guidance(model, x_T, T, alphas, betas, alphas_cumprod,
                            guidance_scale=7.5, condition=None, device='cpu'):
    """
    无分类器引导采样

    参数:
        model: 支持条件生成的模型
        x_T: 初始噪声
        T: 总步数
        alphas, betas, alphas_cumprod: 调度表
        guidance_scale: 引导强度（通常5-15）
        condition: 条件（如文本提示）
        device: 设备

    返回:
        生成的图像
    """
    model.eval()
    x_t = x_T.to(device)

    alphas_cumprod_prev = torch.cat([torch.tensor([1.0]), alphas_cumprod[:-1]])

    with torch.no_grad():
        for t in reversed(range(T)):
            alpha_t = alphas[t]
            beta_t = betas[t]
            alpha_t_bar = alphas_cumprod[t]

            t_tensor = torch.full((x_t.shape[0],), t, device=device, dtype=torch.long)

            # 预测条件噪声和无条件噪声
            noise_cond = model(x_t, t_tensor, condition)
            noise_uncond = model(x_t, t_tensor, None)

            # 组合预测
            predicted_noise = noise_uncond + guidance_scale * (noise_cond - noise_uncond)

            # 更新
            sqrt_recip_alpha_t = 1 / torch.sqrt(alpha_t)
            sqrt_one_minus_alpha_t_bar = torch.sqrt(1 - alpha_t_bar)

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

            if t > 0:
                alpha_t_bar_prev = alphas_cumprod_prev[t]
                posterior_variance = beta_t * (1 - alpha_t_bar_prev) / (1 - alpha_t_bar)
                noise = torch.randn_like(x_t)
                x_t = mean + torch.sqrt(posterior_variance) * noise
            else:
                x_t = mean

    return x_t

def visualize_samples(samples, nrow=4, title="生成样本"):
    """
    可视化生成的样本

    参数:
        samples: [batch_size, channels, height, width]
        nrow: 每行显示的图片数
        title: 标题
    """
    from torchvision.utils import make_grid

    # 归一化到[0, 1]
    samples = (samples - samples.min()) / (samples.max() - samples.min())

    # 创建网格
    grid = make_grid(samples, nrow=nrow, padding=2, normalize=False)

    # 显示
    plt.figure(figsize=(12, 12))
    plt.imshow(grid.permute(1, 2, 0).cpu())
    plt.title(title)
    plt.axis('off')
    plt.show()

# 生成并可视化
samples = torch.randn(16, 3, 32, 32)  # 模拟生成结果
visualize_samples(samples, nrow=4, title="DDPM生成样本")

def calculate_fid(real_images, generated_images, batch_size=32):
    """
    计算FID分数

    参数:
        real_images: 真实图像 [N, C, H, W]
        generated_images: 生成图像 [N, C, H, W]
        batch_size: 批大小

    返回:
        FID分数（越低越好）
    """
    from torchvision.models import inception_v3, Inception_V3_Weights
    from scipy import linalg

    # 加载预训练的Inception模型，移除最后分类层以获取 2048 维 pool 特征
    inception = inception_v3(weights=Inception_V3_Weights.DEFAULT, transform_input=False)
    inception.fc = torch.nn.Identity()  # 替换分类层，输出 2048 维特征
    inception = inception.eval()

    def get_features(images):
        """提取Inception特征"""
        features = []
        with torch.no_grad():
            for i in range(0, len(images), batch_size):
                batch = images[i:i+batch_size]
                # 调整大小到299x299（Inception的输入大小）
                batch_resized = torch.nn.functional.interpolate(
                    batch, size=(299, 299), mode='bilinear', align_corners=False
                )
                feat = inception(batch_resized)
                features.append(feat)
        return torch.cat(features, dim=0).cpu().numpy()

    # 提取特征
    real_features = get_features(real_images)
    gen_features = get_features(generated_images)

    # 计算均值和协方差
    mu_real = np.mean(real_features, axis=0)
    mu_gen = np.mean(gen_features, axis=0)
    sigma_real = np.cov(real_features, rowvar=False)
    sigma_gen = np.cov(gen_features, rowvar=False)

    # 计算FID
    diff = mu_real - mu_gen
    covmean, _ = linalg.sqrtm(sigma_real @ sigma_gen, disp=False)

    if np.iscomplexobj(covmean):
        covmean = covmean.real

    fid = diff @ diff + np.trace(sigma_real + sigma_gen - 2 * covmean)

    return fid

class DiffusionSampler:
    """扩散模型采样器"""

    def __init__(self, model, T=1000, beta_start=0.0001, beta_end=0.02, device='cpu'):
        """
        初始化采样器

        参数:
            model: 训练好的模型
            T: 总步数
            beta_start, beta_end: beta范围
            device: 设备
        """
        self.model = model
        self.T = T
        self.device = device

        # 创建调度表
        self.alphas, self.betas, self.alphas_cumprod = self._get_schedule(T, beta_start, beta_end)
        self.alphas_cumprod_prev = torch.cat([torch.tensor([1.0]), self.alphas_cumprod[:-1]])

    def _get_schedule(self, T, beta_start, beta_end):
        """创建噪声调度表"""
        betas = torch.linspace(beta_start, beta_end, T)
        alphas = 1 - betas
        alphas_cumprod = torch.cumprod(alphas, dim=0)
        return alphas, betas, alphas_cumprod

    def sample(self, batch_size=1, num_steps=None, method='ddpm', **kwargs):  # *args接收任意位置参数，**kwargs接收任意关键字参数
        """
        采样

        参数:
            batch_size: 批大小
            num_steps: 采样步数（None表示使用T）
            method: 采样方法 ('ddpm' 或 'ddim')
            **kwargs: 其他参数

        返回:
            生成的图像 [batch_size, channels, height, width]
        """
        # 创建初始噪声
        x_T = torch.randn(batch_size, 3, 32, 32)  # 假设输入大小为32x32

        if method == 'ddpm':
            return self._ddpm_sample(x_T, num_steps or self.T)
        elif method == 'ddim':
            return self._ddim_sample(x_T, num_steps or self.T)
        else:
            raise ValueError(f"未知的采样方法: {method}")

    def _ddpm_sample(self, x_T, num_steps):
        """DDPM采样"""
        self.model.eval()
        x_t = x_T.to(self.device)

        # 计算子采样间隔
        if num_steps < self.T:
            step_indices = torch.linspace(0, self.T-1, num_steps, dtype=torch.long)
        else:
            step_indices = torch.arange(self.T)

        with torch.no_grad():
            for i, t in enumerate(reversed(step_indices)):
                t = t.item()
                alpha_t = self.alphas[t]
                beta_t = self.betas[t]
                alpha_t_bar = self.alphas_cumprod[t]
                alpha_t_bar_prev = self.alphas_cumprod_prev[t]

                # 预测噪声
                t_tensor = torch.full((x_t.shape[0],), t, device=self.device, dtype=torch.long)
                predicted_noise = self.model(x_t, t_tensor)

                # 计算均值
                sqrt_recip_alpha_t = 1 / torch.sqrt(alpha_t)
                sqrt_one_minus_alpha_t_bar = torch.sqrt(1 - alpha_t_bar)

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

                # 添加噪声
                if i < len(step_indices) - 1:
                    posterior_variance = beta_t * (1 - alpha_t_bar_prev) / (1 - alpha_t_bar)
                    noise = torch.randn_like(x_t)
                    x_t = mean + torch.sqrt(posterior_variance) * noise
                else:
                    x_t = mean

        return x_t

    def _ddim_sample(self, x_T, num_steps):
        """DDIM采样"""
        self.model.eval()
        x_t = x_T.to(self.device)

        # 选择采样时间步
        step_indices = torch.linspace(0, self.T-1, num_steps, dtype=torch.long)

        with torch.no_grad():
            for i, t in enumerate(reversed(step_indices)):
                t = t.item()
                alpha_t = self.alphas[t]
                alpha_t_bar = self.alphas_cumprod[t]

                # 预测噪声
                t_tensor = torch.full((x_t.shape[0],), t, device=self.device, dtype=torch.long)
                predicted_noise = self.model(x_t, t_tensor)

                # 计算x_0的预测
                sqrt_alpha_t_bar = torch.sqrt(alpha_t_bar)
                sqrt_one_minus_alpha_t_bar = torch.sqrt(1 - alpha_t_bar)
                x_0_pred = (x_t - sqrt_one_minus_alpha_t_bar * predicted_noise) / sqrt_alpha_t_bar

                # 计算x_{t-1}
                if i < len(step_indices) - 1:
                    t_prev = step_indices[len(step_indices) - 2 - i].item()
                    alpha_t_prev_bar = self.alphas_cumprod[t_prev]

                    sqrt_alpha_t_prev_bar = torch.sqrt(alpha_t_prev_bar)
                    sqrt_one_minus_alpha_t_prev_bar = torch.sqrt(1 - alpha_t_prev_bar)

                    x_t = sqrt_alpha_t_prev_bar * x_0_pred + sqrt_one_minus_alpha_t_prev_bar * predicted_noise
                else:
                    x_t = x_0_pred

        return x_t

# 使用示例
sampler = DiffusionSampler(model, T=1000, device=device)

# DDPM采样
samples_ddpm = sampler.sample(batch_size=4, num_steps=1000, method='ddpm')
print(f"DDPM采样完成，形状: {samples_ddpm.shape}")

# DDIM采样（更快）
samples_ddim = sampler.sample(batch_size=4, num_steps=50, method='ddim')
print(f"DDIM采样完成，形状: {samples_ddim.shape}")