Denoising Diffusion Probabilistic Models
参考理解DDPM
生成式扩散模型详解：从DDPM到GLIDE
关于Stable diffusion的各种模型，看这篇就可以了！
生成对抗网络（GANs）、自回归模型(AR)、流(FLows)和变分自编码器（VAEs）已经合成了引人注目的图像和音频样本，基于能量的建模和分数匹配取得了显著进展，产生了与GANs相当的图像。
扩散概率模型（为简洁起见，我们称之为“扩散模型”）是一个参数化的马尔可夫链，使用变分推理训练，在有限时间后产生与数据匹配的样本。学习该链的转换以逆转扩散过程，这是一个马尔可夫链，它在采样的相反方向上逐渐向数据中添加噪声，直到信号被破坏。
代码阅读理解
DDPM-UNet代码详解
1.选择数据
选一个数据集，本例采用sklearn自带数据集：
import matplotlib.pyplot as plt
import numpy as np
from sklearn.datasets import make_s_curve
import torch
# TODO 实验数据
s_curve , _ = make_s_curve(10**4 , noise = 0.1)
s_curve = s_curve[:,[0,2] ]/10.0
print("shape of moons :",np.shape(s_curve))
plt.scatter(s_curve[:,0],s_curve[:,1], color='red',edgecolor='white')
plt.savefig('s_curve.png')
2.计算参数
计算前向过程中需要的常数：
# 准备好alpha(αt), beta(βt), αt^, 根号下αt^, 根号下(1-αt^)等值
num_steps = 100 # 设置步长
# 制定每一步的beta
betas = torch.linspace(-6, 6, num_steps) # 逐渐递增
betas = torch.sigmoid(betas)*(0.5e-2 - 1e-5)+1e-5 # β0,β1,...,βt
# 计算alpha、alpha_prod、alpha_prod_previous、alpha_bar_sqrt等变量的值
alphas = 1 - betas # αt = 1 - βt
alphas_prod = torch.cumprod(alphas,0) # αt^ = αt的累乘
# αt^往右平移一位, 原第t步的值维第t-1步的值, 第0步补1
alphas_prod_p = torch.cat([torch.tensor([1]).float(), alphas_prod[:-1]], 0) # αt-1 ^
alphas_bar_sqrt = torch.sqrt(alphas_prod) # αt^ 开根号
one_minus_alphas_bar_log = torch.log(1 - alphas_prod) # log(1-αt^)
one_minus_alphas_bar_sqrt = torch.sqrt(1 - alphas_prod) # 根号下(1-αt^)
assert alphas.shape==alphas_prod.shape==alphas_prod_p.shape==alphas_bar_sqrt.shape==one_minus_alphas_bar_log.shape==one_minus_alphas_bar_sqrt.shape
print("all the same shape",betas.shape)
3.前向计算
# 计算任意时刻的x采样值，基于x_0和重参数化
def q_x(x_0, t):
"""
作用：前向过程, 可以基于x[0]得到任意时刻t的x[t]
输入：x_0:初始干净图像；t：采样步
输出：x_t:第t步时的x_0已成为的样子
"""
noise = torch.randn_like(x_0) # noise为从正态分布中采样的随机噪声
alphas_t = alphas_bar_sqrt[t] # 根号下αt^
alphas_1_m_t = one_minus_alphas_bar_sqrt[t] # 根号下(1-αt^)
return (alphas_t * x_0 + alphas_1_m_t * noise) # 在x[0]的基础上添加噪声
4.前向加噪
演示原始数据分布加噪100步后的结果，可观察到从S曲线在100步中逐渐变为高斯分布的过程：
num_shows = 20
fig,axs = plt.subplots(2, 10, figsize=(28,3))
plt.rc('text',color='black')
#共有10000个点，每个点包含两个坐标
#生成100步以内每隔5步加噪声后的图像
for i in range(num_shows):
j = i//10
k = i%10
q_i = q_x(dataset, torch.tensor([i*num_steps//num_shows])) # 生成t时刻的采样数据
axs[j,k].scatter(q_i[:,0], q_i[:,1], color='green', edgecolor='white')
axs[j,k].set_axis_off()
axs[j,k].set_title('$q(\mathbf{x}_{'+str(i*num_steps//num_shows)+'})$')
很明显每一步我们都在加噪声，这就是正向过程，直至最终是个高斯噪声。
5.扩散网络
写一个简单的网路，用于预测噪声
import torch
import torch.nn as nn
class MLPDiffusion(nn.Module):
def __init__(self, n_steps, num_units=128):
super(MLPDiffusion,self).__init__()
self.linears = nn.ModuleList(
[
nn.Linear(2, num_units),
nn.ReLU(),
nn.Linear(num_units,num_units),
nn.ReLU(),
nn.Linear(num_units,num_units),
nn.ReLU(),
nn.Linear(num_units,2),
]
)
self.step_embeddings = nn.ModuleList(
[
nn.Embedding(n_steps,num_units), # [100,128]
nn.Embedding(n_steps,num_units),
nn.Embedding(n_steps,num_units),
]
)
def forward(self, x, t):
for idx, embedding_layer in enumerate(self.step_embeddings):
t_embedding = embedding_layer(t) # 选第t步的Embedding
x = self.linears[2*idx](x) # 先送入Linear层
x += t_embedding # 加上Embedding
x = self.linears[2*idx+1](x) # 再送入ReLU层
x = self.linears[-1](x) # 最后一个Linear层, 输出为[10000, 2]
return x
6.损失函数
编写训练的误差函数，计算网络预测噪声与真实添加噪声的误差：
def diffusion_loss_fn(model, x_0, alphas_bar_sqrt, one_minus_alphas_bar_sqrt, n_steps):
"""
作用：对任意时刻t进行采样计算loss
参数：
model: 模型
x_0: 干净的图
alphas_bar_sqrt: 根号下αt^
one_minus_alphas_bar_sqrt: 根号下(1-αt^)
n_steps: 采样步
"""
batch_size = x_0.shape[0]
# 对一个batchsize样本生成随机的时刻t, 覆盖到更多不同的t
t = torch.randint(0, n_steps, size=(batch_size//2,)) # 在0~99内生成整数采样步
t = torch.cat([t, n_steps-1-t], dim=0) # 一个batch的采样步, 尽量让生成的t不重复
t = t.unsqueeze(-1) # 增加一个维度(8,1)
# x0的系数
a = alphas_bar_sqrt[t] # 根号下αt^
# eps的系数
aml = one_minus_alphas_bar_sqrt[t] # 根号下(1-αt^)
# 生成随机噪音eps
e = torch.randn_like(x_0)
# 构造模型的输入
x = x_0*a+e*aml # 前向过程：根号下αt^ * x0 + 根号下(1-αt^) * eps
# 送入模型，得到t时刻的随机噪声预测值
output = model(x, t.squeeze(-1)) # 模型预测的是噪声, 噪声维度与x0一样大, [10000,2]
# 与真实噪声一起计算误差，求平均值
return (e - output).square().mean()
7.逆向过程
编写逆扩散采样函数，从随机噪声生成样本
def p_sample_loop(model, shape, n_steps, betas, one_minus_alphas_bar_sqrt):
"""
作用：从x[T]恢复x[T-1]、x[T-2]、...x[0]
输入：
model:模型
shape:数据大小,用于生成随机噪声
n_steps:逆扩散总步长
betas: βt
one_minus_alphas_bar_sqrt: 根号下(1-αt^)
输出：
x_seq：一个序列的x, 即 x[T]、x[T-1]、x[T-2]、...x[0]
"""
cur_x = torch.randn(shape) # 随机噪声, 对应xt
x_seq = [cur_x]
for i in reversed(range(n_steps)):
cur_x = p_sample(model, cur_x, i, betas, one_minus_alphas_bar_sqrt)
x_seq.append(cur_x)
return x_seq
def p_sample(model, x, t, betas, one_minus_alphas_bar_sqrt):
"""
作用：从x[T]采样t时刻的重构值
输入：
model:模型
x: 采样的随机噪声x[T]
t: 采样步
betas: βt
one_minus_alphas_bar_sqrt: 根号下(1-αt^)
输出：
sample: 样本
"""
t = torch.tensor([t])
coeff = betas[t] / one_minus_alphas_bar_sqrt[t] # 模型输出的系数：βt/根号下(1-αt^) = 1-αt/根号下(1-αt^)
eps_theta = model(x, t) # 模型的输出: εθ(xt, t)
# (1/根号下αt) * (xt - (1-αt/根号下(1-αt^))*εθ(xt, t))
mean = (1/(1-betas[t]).sqrt())*(x-(coeff*eps_theta))
z = torch.randn_like(x) # 对应公式中的 z
sigma_t = betas[t].sqrt() # 对应公式中的 σt
sample = mean + sigma_t * z
return (sample)
8.模型训练
开始训练模型，打印loss及中间重构效果
print('Training model...')
batch_size = 128
dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)
num_epoch = 4000
plt.rc('text',color='blue')
model = MLPDiffusion(num_steps) # 输出维度是2，输入是x和step
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
for t in range(num_epoch):
for idx, batch_x in enumerate(dataloader):
# 损失计算
loss = diffusion_loss_fn(model, batch_x, alphas_bar_sqrt, one_minus_alphas_bar_sqrt, num_steps)
optimizer.zero_grad() # 梯度清零
loss.backward() # 损失回传
torch.nn.utils.clip_grad_norm_(model.parameters(),1.) # 梯度裁剪
optimizer.step()
if(t % 100 == 0):
print(loss)
x_seq = p_sample_loop(model, dataset.shape, num_steps, betas, one_minus_alphas_bar_sqrt)
fig, axs = plt.subplots(1, 10, figsize=(28,3))
for i in range(1, 11):
cur_x = x_seq[i*10].detach()
axs[i-1].scatter(cur_x[:,0],cur_x[:,1],color='red',edgecolor='white');
axs[i-1].set_axis_off();
axs[i-1].set_title('$q(\mathbf{x}_{'+str(i*10)+'})$')
扫一扫
694
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
