Diffusion模型

参数说明

beta_schedule = np.linspace(0.0001, 0.02, 1000)  # 示例 beta schedule

alpha_hat = np.cumprod(1 - beta_schedule)  # 计算 alpha_hat

具体例子

让我们通过一个具体的例子展示如何计算 Beta Schedule 和 Alpha_hat：

import numpy as np

# 定义 Beta Schedule
beta_schedule = np.linspace(0.1, 0.2, 10)  # 从0.1到0.2等间隔分成10个值

# 计算 Alpha_hat
alpha_hat = np.cumprod(1 - beta_schedule)

# 打印结果
print("Beta schedule:", beta_schedule)
print("Alpha_hat:", alpha_hat)

在代码中的应用

def encode(x_0, t, alpha_hat):
    alpha_hat_t = alpha_hat[t]  # 获取时间步 t 的 alpha_hat
    epsilon = torch.randn_like(x_0)  # 生成标准正态分布噪声
    epsilon_prime = torch.sign(x_0) * torch.abs(epsilon)  # 计算符号噪声
    x_t = torch.sqrt(alpha_hat_t) * x_0 + torch.sqrt(1 - alpha_hat_t) * epsilon_prime  # 计算时间步 t 的特征向量
    return x_t

for t in range(num_steps):
    x_t = encode(x_0, t, alpha_hat)
    node_features.append(x_t.numpy())
    print(f"Node features at time step {t+1} (x_{t+1}):", x_t.numpy())

反向扩散过程和损失计算

for t in range(num_steps-1, -1, -1):
    x_t = torch.tensor(node_features[t+1])
    mu_theta, sigma_t = model.decode(x_t, t, alpha_hat)
    print(f"Reconstructed features at time step {t} (mu_theta_{t}):", mu_theta.detach().numpy())

完整的代码

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

定义模型类
class DiffusionModel(nn.Module):
    def __init__(self, d):
        super(DiffusionModel, self).__init__()
        self.epsilon_predictor = nn.Sequential(
            nn.Linear(d, d),
            nn.ReLU(),
            nn.Linear(d, d)
        )

    def decode(self, x_t, t, alpha_hat):
        alpha_hat_t = alpha_hat[t]
        mu_theta = self.epsilon_predictor(x_t)
        sigma_t = torch.sqrt(1 - alpha_hat_t)
        return mu_theta, sigma_t

def encode(x_0, t, alpha_hat):
    alpha_hat_t = alpha_hat[t]
    epsilon = torch.randn_like(x_0)
    epsilon_prime = torch.sign(x_0) * torch.abs(epsilon)
    x_t = torch.sqrt(alpha_hat_t) * x_0 + torch.sqrt(1 - alpha_hat_t) * epsilon_prime
    return x_t, epsilon

def compute_total_loss(model, x_0, alpha_hat):
    num_steps = len(alpha_hat)
    x_t = x_0
    epsilon_list = []
    
    多步噪声传递过程
    for t in range(num_steps):
        x_t, epsilon = encode(x_t, t, alpha_hat)
        epsilon_list.append(epsilon)
    
    total_loss = 0
    
    反向扩散过程
    for t in range(num_steps-1, -1, -1):
        mu_theta, sigma_t = model.decode(x_t, t, alpha_hat)
        epsilon_t = epsilon_list[t]
        total_loss += F.mse_loss(mu_theta, epsilon_t)
        更新 x_t 为上一步的特征
        x_t = (x_t - torch.sqrt(1 - alpha_hat[t]) * epsilon_t) / torch.sqrt(alpha_hat[t])
    
    return total_loss

初始化参数
d = 5  # 特征向量维度
beta_schedule = np.linspace(0.1, 0.2, 10)  # 示例 beta 计划
alpha_hat = np.cumprod(1 - beta_schedule)  # 计算 alpha_hat

初始化节点特征向量
x_0 = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0])
print("Initial node features (x_0):", x_0.numpy())

初始化模型
model = DiffusionModel(d)

计算总损失
total_loss = compute_total_loss(model, x_0, alpha_hat)
print("Total loss:", total_loss.item())