小白读论文——Denoising Diffusion Probabilistic Model 以及代码理解

bin.0327

已于 2024-08-26 11:34:00 修改

阅读量927

点赞数 23

文章标签：计算机视觉人工智能

于 2024-08-22 10:18:22 首次发布

本文链接：https://blog.csdn.net/weixin_74479558/article/details/141377107

版权

小白读论文——Denoising Diffusion Probabilistic Model 以及代码理解

一、概述 ——石头里的雕塑
二、过程解析
三、基础知识
四、公式推导
五、代码理解
扩散过程
去噪过程

学习视频：链接: b站博主：Enzo_Mi

一、概述 ——石头里的雕塑

在这里插入图片描述

二、过程解析

1.加噪过程

下一时刻的图像是由上一时刻的图像加上噪音得到的
加噪过程
$x_t=\sqrt{\beta_t} \times \epsilon_t + \sqrt{1-\beta_t} \times x_{t-1}$
$x_t$ : $t$ 时刻的图像
$\epsilon_t$ : $t$ 时刻随机生成的噪声图像
$\beta_t$ :是提前准备好的常数，不同时刻的值不同，随着 $t$ 的增大而增大，因此表明越到后面，加的噪音越多

因此，根据公式可以看出，下一时刻的加噪图像是由两部分组成，一部分是噪声图像，一部分是上一时刻的图像，分别乘以权重相加得到

任意时刻的 $x_t$ 可由原图像和噪声生成

$x_t =\sqrt{1-\bar{\alpha_t}} \times \epsilon +\sqrt{\bar{\alpha_t}} \times x_0 , \alpha_t=1-\beta_t, \bar{\alpha}=\alpha_t\alpha_{t-1}......\alpha_2 \alpha_1$

2. training

在这里插入图片描述
重复此过程直到收敛
2：从数据集取出一张图片
3：从1到T中取出一个 $t$ 值
4：取出一个 $\epsilon$
5：使用梯度下降法优化损失函数 $\theta$ 的值， $\theta$ 是Unet网络中所有参数的集合
整个流程就是
首先生成 $x_t$ ，然后将 $x_t$ 和 $t$ 输入Unet网络中，输出预测的噪声 $\epsilon$ , 最后与原始真实噪声比较得到误差

3. sampling

在这里插入图片描述
1：先从随机分布中采样出一个 $x_T$
2：要循环T次，每一次生成前一张图像，直到生成到最原始的图像
3：从随机分布中采样出 $z$
4：逆推上一时刻的图像 $x_{t-1}$

三、基础知识

在这里插入图片描述

四、公式推导

1.加噪过程

要推到的两个公式
在这里插入图片描述

2.去噪过程

2个假设：
(1) 假设去噪过程也是一个马尔可夫链（前一时刻的图像只与后一时刻的图像有关）
(2) 假设去噪过程是高斯分布
在这里插入图片描述

3.损失函数

极大似然估计的思想
（1）假设图像（每一个像素）服从正态分布
（2）我们要找出一对均值和方差，使得从该正态分布中采样出这些样本值的概率最大
在这里插入图片描述

最终目标函数为

$q(x_{t-1}|x_t,x_0)$ :已知 $x_t,x_0$ 推导出 $x_{t-1}$
$p_\theta$ :没法计算，需要用神经网络（Unet）进行预测
最终使得 $p_\theta$ 的分布越来越拟合于 $q$

五、代码理解

参考视频： b站博主：deep_thoughts
参考博客: csdn:一文弄懂 Diffusion Model（DDPM）+ 代码实现
参考代码代码

扩散过程

1.生成数据集

使用sklearn生成一个10000个点，每个点的维度是2的数据集

%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
from sklearn.datasets import make_s_curve # 生成S形二维数据点 https://scikit-learn.org/stable/modules/generated/sklearn.datasets.make_s_curve.html
import torch

s_curve,_ = make_s_curve(10**4,noise=0.1) # 10000个数据点
s_curve = s_curve[:,[0,2]]/10.0

print("shape of s:",np.shape(s_curve))

data = s_curve.T

fig,ax = plt.subplots()
ax.scatter(*data,color='blue',edgecolor='white');

ax.axis('off')

dataset = torch.Tensor(s_curve).float()

在这里插入图片描述

2.定义超参数的值

betas: $\beta$

alphas: $\alpha = 1-\beta$

alphas_prod: $\overline{\alpha_t} = \prod_{s=1}^{t}\alpha_s$

alphas_prod_p: $alphas\_prod[:-1] = \overline{\alpha_{t-1}}$

alphas_bar_sqrt: $\sqrt{\overline{\alpha_t}}$

one_minus_alphas_bar_log: $log(1-\overline{\alpha_t})$

one_minus_alphas_bar_sqrt: $\sqrt{1-\overline{\alpha_t}}$

num_steps = 100

#制定每一步的beta
betas = torch.linspace(-6,6,num_steps)
betas = torch.sigmoid(betas)*(0.5e-2 - 1e-5)+1e-5

#计算alpha、alpha_prod、alpha_prod_previous、alpha_bar_sqrt等变量的值
alphas = 1-betas
alphas_prod = torch.cumprod(alphas,0)
alphas_prod_p = torch.cat([torch.tensor([1]).float(),alphas_prod[:-1]],0)
alphas_bar_sqrt = torch.sqrt(alphas_prod)
one_minus_alphas_bar_log = torch.log(1 - alphas_prod)
one_minus_alphas_bar_sqrt = torch.sqrt(1 - alphas_prod)

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_t = \sqrt{\overline{\alpha_t}}x_0+\sqrt{1-\overline{\alpha_t}}\epsilon$

#计算任意时刻的x采样值，基于x_0和重参数化
def q_x(x_0,t):
    """可以基于x[0]得到任意时刻t的x[t]"""
    # x_0：dataset，原始(10000, 2)的数据点集
    # t: torch.tensor([i]),i为采样几次
    noise = torch.randn_like(x_0)
    alphas_t = alphas_bar_sqrt[t]
    alphas_1_m_t = one_minus_alphas_bar_sqrt[t]
    return (alphas_t * x_0 + alphas_1_m_t * noise)#在x[0]的基础上添加噪声

4.演示原始数据加噪95步后的结果

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='red',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),
                nn.Embedding(n_steps,num_units),
                nn.Embedding(n_steps,num_units),
            ]
        )
    def forward(self,x,t):
        #  x = x_0
        for idx,embedding_layer in enumerate(self.step_embeddings):
            t_embedding = embedding_layer(t)
            x = self.linears[2*idx](x)
            x += t_embedding
            x = self.linears[2*idx+1](x)
            
        x = self.linears[-1](x)
        
        return x
# nn.Embedding：https://blog.csdn.net/qq_39540454/article/details/115215056
# 使用示例：输出维度是2，输入是x和step
# model = MLPDiffusion(num_steps)
# output = model(x,step)

6.编写训练的误差函数

model = MLPDiffusion(num_steps)

x_0: batch_x,部分x_0

alphas_bar_sqrt： $\sqrt{\overline{\alpha_t}}$

one_minus_alphas_bar_sqrt： $\sqrt{1-\overline{\alpha_t}}$

n_steps: num_steps = 100

loss = diffusion_loss_fn(model,x_0,alphas_bar_sqrt,one_minus_alphas_bar_sqrt,n_steps)

def diffusion_loss_fn(model,x_0,alphas_bar_sqrt,one_minus_alphas_bar_sqrt,n_steps):
    """对任意时刻t进行采样计算loss"""
    batch_size = x_0.shape[0]
    
    #对一个batchsize样本生成随机的时刻t,t的形状是torch.Size([batchsize, 1])
    t = torch.randint(0,n_steps,size=(batch_size//2,))
    t = torch.cat([t,n_steps-1-t],dim=0)
    t = t.unsqueeze(-1)
    
    #x0的系数
    a = alphas_bar_sqrt[t] # torch.Size([batchsize, 1])
    
    #eps的系数
    aml = one_minus_alphas_bar_sqrt[t] # torch.Size([batchsize, 1])
    
    #生成随机噪音eps
    e = torch.randn_like(x_0) # torch.Size([batchsize, 2])
    
    #构造模型的输入
    x = x_0*a+e*aml # torch.Size([batchsize, 2])
    
    #送入模型，得到t时刻的随机噪声预测值
    output = model(x,t.squeeze(-1)) #t.squeeze(-1)为torch.Size([batchsize])
    # output:torch.Size([batchsize, 2])
    #与真实噪声一起计算误差，求平均值
    return (e - output).square().mean()

7.训练模型

## ----------------------------- 训练模型 ----------------------------- ##
print('Training model...')
batch_size = 128
num_epoch = 4000
dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)

model = MLPDiffusion(num_steps)
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)
       torch.save(model.state_dict(), 'model_{}.pth'.format(t))

去噪过程

$x_t-->x_{t-1}$

1.定义单步逆扩散过程

def p_sample(model, x, t, betas, one_minus_alphas_bar_sqrt):
    """
    从x[t]采样t-1时刻的重构值x[t-1]
    :param model:
    :param x: x[T]
    :param t:
    :param betas:
    :param one_minus_alphas_bar_sqrt:
    :return:
    """
    ## 1) 求出 bar_u_t
    t = torch.tensor([t])
    coeff = betas[t] / one_minus_alphas_bar_sqrt[t]
    # 送入U-Net模型，得到t时刻的随机噪声预测值 eps_theta
    eps_theta = model(x, t)
    mean = (1 / (1 - betas[t]).sqrt()) * (x - (coeff * eps_theta))
 
    ## 2) 得到 x[t-1]
    z = torch.randn_like(x)
    sigma_t = betas[t].sqrt()
    sample = mean + sigma_t * z
    return (sample)

2.根据单步逆扩散过程定义循环逆扩散过程

def p_sample_loop(model, noise_x_t, n_steps, betas, one_minus_alphas_bar_sqrt):
   """
   从x[T]恢复x[T-1]、x[T-2]|...x[0] 的循环
   :param model:
   :param shape:数据集的形状，也就是x[T]的形状
   :param n_steps:
   :param betas:
   :param one_minus_alphas_bar_sqrt:
   :return: x_seq由x[T]、x[T-1]、x[T-2]|...x[0]组成, cur_x是从噪声中生成的图片
   """
   # 得到噪声x[T]
   cur_x = noise_x_t
   x_seq = [noise_x_t]
   # 从x[T]恢复x[T-1]、x[T-2]|...x[0]
   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, cur_x

3.根据训练好的扩散模型+随机噪声——>生成数据

# 1) 加载训练好的diffusion model
model = MLPDiffusion(num_steps)
model.load_state_dict(torch.load('model_1000.pth'))
# 2) 生成随机噪声x[T]
noise_x_t = torch.randn(dataset.shape)
# 3) 根据随机噪声逆扩散为x[T-1]、x[T-2]|...x[0] + 图片x[0]
x_seq, cur_x = p_sample_loop(model, noise_x_t, num_steps, betas, one_minus_alphas_bar_sqrt)
print(x_seq, cur_x)