Diffusion model

Stable Diffusion------DDMP 代码理解

Stable dissuison model 中的 .yaml文件是模型的配置文件，改配置文件中主要包含model 和 data，今天就浅浅记录model 的相关配置和理解。

model:
  base_learning_rate: 1.0e-04
  target: ldm.models.diffusion.ddpm.LatentDiffusion

ok, 来到 ldm.models.diffusion.ddpm.LatentDiffusion，

class LatentDiffusion(DDPM):

看到了LatentDiffusion 继承了 DDMP类，so，先来看DDMP类。
在看代码之前，肯定要回顾一下 Diffuison model 的前向过程（加噪过程）和后向过程（去噪过程）。

diffuison 前向过程

前向过程就是对一张图像进行加噪声：
X₀ —> X₁ —> X₂ —> X₃—> …—>X_t
X₀ 为清晰的图像， X_t 为加了N步噪声的图像。

此刻的疑问
如果你是第一次了解stable difuison model 你肯定会问，噪声具体是什么？又是怎么加的？
ok，噪声是从正太分布采样获得的，并且噪声是越来越大的（也是符合一个正常思维的，慢慢的加噪声，噪声逐渐递增，直到图像模糊）。

公式推导

怎么加？首先明白加噪声的一个过程：
$$P(X_t|X_{t-1})=N(X_t;\sqrt{1-{\beta }_t}{X}_{t-1};{\beta }_{t}I).$$
论文中和代码中，都令$${\alpha }_t=1-{\beta }_t$$
转化为具体的公式如下：
$$X_t=\sqrt{{\alpha }_t}{X}_{t-1}+\sqrt{1-{\alpha }_t}{z}_1（z是服从正太分布的，在这里看作是一个噪声）$$
那么
$$X_{t-1}=\sqrt{{\alpha }_{t-1}}{X}_{t-2}+\sqrt{1-{\alpha }_{t-1}}{z}_2$$
上面两个公式合并，最终可以推出：
$$X_t=\sqrt{{\bar{\alpha }}_t}{X}_0+\sqrt{1-{\bar{\alpha }}_t}z$$
具体推出的过程如下：

代码理解

ook，终于到了对应的代码部分，DDMP类中的一个函数：

    def register_schedule(self, given_betas=None, beta_schedule="linear", timesteps=1000,
                          linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
        if exists(given_betas):
            betas = given_betas
        else:
            betas = make_beta_schedule(beta_schedule, timesteps, linear_start=linear_start, linear_end=linear_end,
                                       cosine_s=cosine_s)
        alphas = 1. - betas
        alphas_cumprod = np.cumprod(alphas, axis=0)
        alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])

        timesteps, = betas.shape
        self.num_timesteps = int(timesteps)
        self.linear_start = linear_start
        self.linear_end = linear_end
        assert alphas_cumprod.shape[0] == self.num_timesteps, 'alphas have to be defined for each timestep'

        to_torch = partial(torch.tensor, dtype=torch.float32)

        self.register_buffer('betas', to_torch(betas))
        self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
        self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev))

改段代码对应前向过程中的参数。
其中self.sqrt_alphas_cumprod​和self.sqrt_one_minus_alphas_cumprod​分别是alpha的累乘值和1-alpha的累乘值。
在前向训练的过程中，噪声是怎么加上的呢？ok，下面来看代码，

    def forward(self, x, *args, **kwargs):
        # b, c, h, w, device, img_size, = *x.shape, x.device, self.image_size
        # assert h == img_size and w == img_size, f'height and width of image must be {img_size}'
        t = torch.randint(0, self.num_timesteps, (x.shape[0],), device=self.device).long()
        return self.p_losses(x, t, *args, **kwargs)

在训练的时候调用的是 **self.p_losses(x, t, *args, kwargs)

    def p_losses(self, x_start, t, noise=None):
        noise = default(noise, lambda: torch.randn_like(x_start))
        x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
        model_out = self.model(x_noisy, t)

        loss_dict = {}
        if self.parameterization == "eps":
            target = noise
        elif self.parameterization == "x0":
            target = x_start
        else:
            raise NotImplementedError(f"Paramterization {self.parameterization} not yet supported")

        loss = self.get_loss(model_out, target, mean=False).mean(dim=[1, 2, 3])

        log_prefix = 'train' if self.training else 'val'

        loss_dict.update({f'{log_prefix}/loss_simple': loss.mean()})
        loss_simple = loss.mean() * self.l_simple_weight

        loss_vlb = (self.lvlb_weights[t] * loss).mean()
        loss_dict.update({f'{log_prefix}/loss_vlb': loss_vlb})

        loss = loss_simple + self.original_elbo_weight * loss_vlb

        loss_dict.update({f'{log_prefix}/loss': loss})

        return loss, loss_dict

noise = default(noise, lambda: torch.randn_like(x_start))
torch.randn_like()函数生成一个与x_start具有相同形状的随机张量。torch.randn_like()函数会生成一个服从标准正态分布（均值为0，方差为1）的随机张量。对应 Z1，Z2。
x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
self.q_sample（）函数是 噪声采样 。即给初始图像加噪声，对应下面的
$$X_t=\sqrt{{\bar{\alpha }}_t}{X}_0+\sqrt{1-{\bar{\alpha }}_t}z$$

    def q_sample(self, x_start, t, noise=None):
        noise = default(noise, lambda: torch.randn_like(x_start))
        return (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
                extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise)

到这里，前向的过程已经结束了。ook，接下来看反向过程即去噪过程。

diffuison 反向过程

公式推导

目标：由X_t 推出 X₀
X₀ <— X₁ <— X₂ <— X₃<— …<—X_t
怎么推？贝叶斯公式！！！！！
回忆贝叶斯公式：

由贝叶斯公式，可以从X_t推出X_t-1
$$q（X_{t-1}|X_t）=\frac{q（X_t|X_t-1）\ast q(x_{t-1})}{q(X_t)}$$
进一步地，加X₀，
$$q（X_{t-1}|X_t,X_0）=\frac{q（X_t|X_t-1,X_0）\ast q(X_{t-1},X_0)}{q(X_t,X_0)}$$
由前向过程中最终的公式
$$X_t=\sqrt{{\bar{\alpha }}_t}{X}_0+\sqrt{1-{\bar{\alpha }}_t}z$$
以及
$$X_t=\sqrt{{\alpha }_t}{X}_{t-1}+\sqrt{1-{\alpha }_t}{z}_1$$
可得出

ook，进一步地，3个分布的乘除到底该怎么算？肯定是用分布的概率密度函数啦，

可写成

对括号里的多项式进行展开，合并整理，在整理得过程中，不要迷了，我们得目的是从X_t推出X_t-1.
所以我们最终求得的是X_t-1的分布：

通过重参数技巧， 最终的方差是常数：
$${\sigma }^2=\frac{1-{\bar{\alpha }}_{t-1}{\beta }_t}{1-{\bar{\alpha }}_t}$$
均值：
$$\mu =\frac{{\sqrt{\alpha }}_t(1-{\bar{\alpha }}_{t-1})X_t}{1-{\bar{\alpha }}_t}+\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t{X}_0}{1-{\bar{\alpha }}_t}$$
其中，X0在反向工程中是未知！！！是要最终求的！！！
由公式
$$X_t=\sqrt{{\bar{\alpha }}_t}{X}_0+\sqrt{1-{\bar{\alpha }}_t}z$$
可以用X_t 表示X₀ ， 并求得最终得均值

$$\mu =\frac{1}{{\sqrt{\alpha }}_t}(X_t-\frac{{\beta }_t{Z}_t}{\sqrt{1-{\bar{\alpha }}_t}})$$
其中，Zt是噪声，由深度网络预测！！！！
反向的过程可以总结为：
$$P(X_{t-1}|X_t,X_0)=N(X_{t-1};\mu ;{\sigma }^{2}I).$$
我们让学习来的 P(X_{t-1}|X_{t},X_{0}) 来逼近 q(X_{t-1}|X_{t}) .
终于完了！

反向代码

首先看 def p_mean_variance(self, x, t, clip_denoised: bool): 函数

     def p_mean_variance(self, x, t, clip_denoised: bool):
        model_out = self.model(x, t)
        if self.parameterization == "eps":
        # 模型预测的是噪声
            x_recon = self.predict_start_from_noise(x, t=t, noise=model_out)
        elif self.parameterization == "x0":
        # 模型预测的是去噪图像
            x_recon = model_out
        if clip_denoised:
            x_recon.clamp_(-1., 1.)

        model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
        return model_mean, posterior_variance, posterior_log_variance

该函数的输入：
x：是当前时间步t的噪声图像Xt
t：表示时间步
clip_denoised：是否将图像的值裁剪到 ( − 1.0 , 1.0 ) 区间.
逐行来看代码：

model_out = self.model(x, t)

model_out是模型预测出的噪声

  x_recon = self.predict_start_from_noise(x, t=t, noise=model_out)

    def predict_start_from_noise(self, x_t, t, noise):
        return (
                extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
                extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise
        )

函数 def predict_start_from_noise(self, x_t, t, noise):
该函数的输入是：
x_t: 带噪音的图像X_t
t: 时间步长
noise: 预测的噪声
输出：
预测的原始图像(是由下列公式推导出)

$$X_t=\sqrt{{\bar{\alpha }}_t}{X}_0+\sqrt{1-{\bar{\alpha }}_t}z$$

$${\tilde{X}}_0=\sqrt{\frac{1}{{\bar{\alpha }}_t}}{X}_t+\sqrt{\frac{1}{{\bar{\alpha }}_t}}-1z$$
到这里感觉是不是 就是到此结束了，现在 直接把未加噪的图像预测出来了吗 ？？？
但是，仔细去想， 通过加了1000步噪声，能直接从1000步加噪后的图像直接预测出原始图像吗？公式是合理的，但是神经网络应该还不够这么强大！！！
继续往下看：

model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start = x_start, x_t = x, t = t)

def q_posterior(self, x_start, x_t, t):
    # 计算平均值
    posterior_mean = (
        extract(self.posterior_mean_coef1, t, x_t.shape) * x_start +
        extract(self.posterior_mean_coef2, t, x_t.shape) * x_t
    )
    # 计算方差
    posterior_variance = extract(self.posterior_variance, t, x_t.shape)
    # 获得一个压缩范围的方差,且取对数
    posterior_log_variance_clipped = extract(self.posterior_log_variance_clipped, t, x_t.shape)
    return posterior_mean, posterior_variance, posterior_log_variance_clipped

输入：
预测的初始图像
x_t: t步长的噪声图像
t:时间步长

输出：
预测的方差
预测的均值

$${\sigma }^2=\frac{1-{\bar{\alpha }}_{t-1}{\beta }_t}{1-{\bar{\alpha }}_t}$$
均值：
$$\mu =\frac{{\sqrt{\alpha }}_t(1-{\bar{\alpha }}_{t-1})X_t}{1-{\bar{\alpha }}_t}+\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t{X}_0}{1-{\bar{\alpha }}_t}$$
其中self.posterior_mean_coef1​对应的是x0前面的系数，self.posterior_mean_coef2​对应的是xt前面的系数。

既然得到了 当前时间步t 预测的方差和均值，就可以采样出下一个时间步长t-1的样本

@torch.no_grad()
def p_sample(self, x, t, clip_denoised=True, repeat_noise=False):
    b, *_, device = *x.shape, x.device
    model_mean, _, model_log_variance = self.p_mean_variance(x=x, t=t, clip_denoised=clip_denoised)
    noise = noise_like(x.shape, device, repeat_noise)  # 返回一个和x一样形状的标准高斯噪音noise
    # repeat_noise表示是否重复使用一个噪音, 若重复使用, 一个batch内的所有样本将加同一个随机噪音; 否则每个样本将独立采样
    # nonzero_mask表示是否有噪音, t=0时无噪音(为0), 其它时候有噪音(为1)
    nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
    return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise

输入：
x： 当前的时间样本
t：时间步长
clip_denoised. 是否对噪音裁剪到区间 ( − 1.0 , 1.0 ) 内.
repeat_noise. 是否在一个批量中对所有样本重复使用同一个噪音

输出：
和 x 相同形状的, 下一个时间步中的一个批量的样本 X_t-1
ook, 继续往下看,

@torch.no_grad()
def p_sample_loop(self, shape, return_intermediates=False):
    device = self.betas.device
    b = shape[0]
    img = torch.randn(shape, device=device)
    intermediates = [img]
    for i in tqdm(reversed(range(0, self.num_timesteps)), desc='Sampling t', total=self.num_timesteps):
        # i从T-1到0
        # t = torch.full((b,), i, device=device, dtype=torch.long)
        img = self.p_sample(img, torch.full((b,), i, device=device, dtype=torch.long),
                            clip_denoised=self.clip_denoised)
        # img是x_i
        if i % self.log_every_t == 0 or i == self.num_timesteps - 1:
            intermediates.append(img)  # 存中间的图像
    if return_intermediates:
        return img, intermediates
    return img  # x_0

输入:

shape. 图像的形状.
return_intermediates. 是否返回反向过程中的中间图像.、

输出:

img. 生成的图像 x 0
intermediates. 一个列表, 存了中间图像.