Spaces:
Runtime error
Runtime error
File size: 7,987 Bytes
de2377a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 | # motion_generation/lib/model/ddpm.py
import torch
import torch.nn as nn
from lib.model.unet1d import UNet1D
class DDPM(nn.Module):
def __init__(self, model: UNet1D, beta_schedule='linear', timesteps=1000):
super().__init__()
self.model = model # 噪声预测器 (UNet1D)
self.timesteps = timesteps
# ----------------------------------------------------
# 1. 定义 Beta 调度和相关参数
# ----------------------------------------------------
# 1.1. 计算 Betas (噪声水平)
if beta_schedule == 'linear':
betas = self._linear_beta_schedule(timesteps)
else:
raise NotImplementedError(f"不支持的 beta_schedule: {beta_schedule}")
# 将所有参数注册为 buffer (不会被训练，但会保存到 state_dict)
self.register_buffer('betas', betas)
# 1.2. 计算 Alpha 相关的参数
alphas = 1.0 - betas
alphas_cumprod = torch.cumprod(alphas, dim=0) # 累积乘积: \bar{\alpha}_t
alphas_cumprod_prev = torch.cat([torch.tensor([1.0]), alphas_cumprod[:-1]])
# 方便计算的根号形式
sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod)
sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - alphas_cumprod)
sqrt_recip_alphas = torch.sqrt(1.0 / alphas) # 1/sqrt(alpha_t)
# 后验均值计算所需的参数 (反向过程)
posterior_variance = betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)
# 注册 buffers
self.register_buffer('alphas', alphas)
self.register_buffer('alphas_cumprod', alphas_cumprod)
self.register_buffer('sqrt_alphas_cumprod', sqrt_alphas_cumprod)
self.register_buffer('sqrt_one_minus_alphas_cumprod', sqrt_one_minus_alphas_cumprod)
self.register_buffer('sqrt_recip_alphas', sqrt_recip_alphas)
self.register_buffer('posterior_variance', posterior_variance)
# ----------------------------------------------------
def _linear_beta_schedule(self, timesteps, start=0.0001, end=0.02):
"""线性 Beta 调度"""
return torch.linspace(start, end, timesteps)
def get_index_from_t(self, variances, t, x_shape):
"""从 (T,) 维度的张量中，根据时间步 t (B,) 提取对应的方差/系数，并重塑至 (B, 1, 1...)"""
B = t.shape[0]
out = variances.gather(-1, t) # 从 variances 中提取 t 对应的值
# 重塑维度以匹配 x 的形状 (B, C) -> (B, C)
# 对于我们的姿态数据 (B, 72)， C=72，所以只需要重塑到 (B, 1)
return out.reshape(B, *([1] * (len(x_shape) - 1)))
# ----------------------------------------------------
# 2. 前向扩散 (加噪)
# ----------------------------------------------------
def forward_diffusion(self, x_start, t, noise=None):
"""
前向过程：计算 x_t = sqrt(alpha_bar_t) * x_0 + sqrt(1 - alpha_bar_t) * noise
Args:
x_start: 初始数据 x_0 (B, 72)
t: 时间步 (B,)
noise: 用于加噪的噪声，如果为 None 则随机生成
Returns:
x_t: t时刻加噪后的数据 (B, 72)
noise: 实际使用的噪声 (B, 72)
"""
if noise is None:
noise = torch.randn_like(x_start)
# 提取 sqrt(alpha_bar_t) 和 sqrt(1 - alpha_bar_t)
sqrt_alpha_bar_t = self.get_index_from_t(self.sqrt_alphas_cumprod, t, x_start.shape)
sqrt_one_minus_alpha_bar_t = self.get_index_from_t(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape)
# 公式实现
x_t = sqrt_alpha_bar_t * x_start + sqrt_one_minus_alpha_bar_t * noise
return x_t, noise
# ----------------------------------------------------
# 3. 损失计算 (训练过程)
# ----------------------------------------------------
def forward(self, x_start):
"""
训练时的前向传播，用于计算损失。
Args:
x_start: 批量原始数据 (归一化后的姿态) (B, 72)
Returns:
loss: 均方误差 (MSE) 损失
"""
B = x_start.shape[0]
# 1. 随机采样时间步 t
t = torch.randint(0, self.timesteps, (B,), device=x_start.device).long()
# 2. 随机生成噪声
noise = torch.randn_like(x_start)
# 3. 计算 x_t 和 t 时刻的真实噪声
x_t, _ = self.forward_diffusion(x_start, t, noise)
# 4. 噪声预测网络预测噪声
predicted_noise = self.model(x_t, t)
# 5. 计算损失：预测噪声和真实噪声的均方误差 (MSE)
loss = nn.functional.mse_loss(predicted_noise, noise)
return loss
# ----------------------------------------------------
# 4. 采样 (生成过程)
# ----------------------------------------------------
@torch.no_grad()
def sample(self, sample_shape, device, return_intermediates=False, save_interval=None):
"""
逆向过程：从纯噪声开始，迭代去噪 T 步，生成新的姿态数据。
Args:
sample_shape: 要生成的样本形状 (B, 72)
device: 运行设备
return_intermediates: 是否返回中间状态
save_interval: 保存中间状态的间隔（timesteps），如果为None则不保存
Returns:
如果 return_intermediates=False:
x_0: 最终生成的归一化姿态 (B, 72)
如果 return_intermediates=True:
(x_0, intermediates): 最终姿态和中间状态列表
intermediates: List[(timestep, x_t)] 包含时间步和对应的状态
"""
# 从纯噪声开始 x_T ~ N(0, I)
x = torch.randn(sample_shape, device=device)
# 存储中间状态
intermediates = []
if return_intermediates:
# 保存初始噪声状态
intermediates.append((self.timesteps, x.clone().cpu()))
# 从 T-1 步迭代到 0 步
for t in reversed(range(0, self.timesteps)):
t_tensor = torch.full((sample_shape[0],), t, device=device, dtype=torch.long)
# 1. 预测噪声
predicted_noise = self.model(x, t_tensor)
# 2. 提取当前时刻的系数
beta_t = self.get_index_from_t(self.betas, t_tensor, x.shape)
sqrt_one_minus_alpha_bar_t = self.get_index_from_t(self.sqrt_one_minus_alphas_cumprod, t_tensor, x.shape)
sqrt_recip_alpha_t = self.get_index_from_t(self.sqrt_recip_alphas, t_tensor, x.shape)
# 3. 计算均值 (mu_t-1)
# DDPM公式: μ_t = (1/sqrt(alpha_t)) * (x_t - (beta_t/sqrt(1-alpha_bar_t)) * predicted_noise)
mean = sqrt_recip_alpha_t * (x - beta_t * predicted_noise / sqrt_one_minus_alpha_bar_t)
# 4. 添加噪声项
if t > 0:
variance = self.get_index_from_t(self.posterior_variance, t_tensor, x.shape)
noise = torch.randn_like(x)
# x_{t-1} = \mu_{t-1} + \sigma_{t-1} * z
x = mean + torch.sqrt(variance) * noise
else:
x = mean # t=0 时不加噪声，直接取均值作为最终输出
# 保存中间状态
if return_intermediates:
if save_interval is None:
# 保存所有状态
intermediates.append((t, x.clone().cpu()))
elif t % save_interval == 0 or t == 0:
# 按间隔保存
intermediates.append((t, x.clone().cpu()))
# 将输出限制在 [-1, 1] 附近（可选，取决于您的归一化范围）
# x = x.clamp(-1., 1.)
if return_intermediates:
return x, intermediates
return x
|