本文代码节选自 mmdet 中的 DiffusionDet 代码,目前该代码还处于 Development 阶段,所以我博客里写的代码和之后的稳定版本可能稍有不同,不过不用担心,我们只看最关键的部分
DDPM中扩散部分有个参数
β
\beta
β:
q
(
z
t
∣
z
t
−
1
)
:
=
N
(
z
t
;
1
−
β
t
z
t
−
1
,
β
t
I
)
q(z_t | z_{t-1}) := \mathcal{N} (z_{t}; \sqrt{1 - \beta_t} z_{t-1}, \beta_t \bf{I} )
q(zt∣zt−1):=N(zt;1−βtzt−1,βtI)
这就是每次的加噪过程,也可以视为
z
t
−
1
z_{t-1}
zt−1先经过一个缩放,再加一个随机噪声之后,就成了
z
t
z_{t}
zt。 每次加噪声通过一个参数
β
t
\beta_t
βt来控制,这个参数是人为给定的,而不是可学习的,由于:
q
(
z
t
∣
z
0
)
:
=
N
(
z
t
;
α
ˉ
t
z
0
,
(
1
−
α
ˉ
t
)
I
)
q(z_t | z_{0}) := \mathcal{N} (z_{t}; \sqrt{ \bar{\alpha}_t } z_{0}, (1-\bar{\alpha}_t) \bf{I} )
q(zt∣z0):=N(zt;αˉtz0,(1−αˉt)I) 即:
z
t
=
α
ˉ
t
z
0
+
ϵ
1
−
α
ˉ
t
,
w
h
e
r
e
ϵ
∈
N
(
0
,
I
)
z_t = \sqrt{ \bar{\alpha}_t } z_{0} + \epsilon \sqrt{1 - \bar{\alpha}_t}, \ \ where \ \ \epsilon \in \mathcal{N}(0, \bf{I})
zt=αˉtz0+ϵ1−αˉt,whereϵ∈N(0,I)
在给定
z
0
z_{0}
z0 的基础上,
q
(
z
t
∣
z
0
)
q(z_t | z_{0})
q(zt∣z0) 也是一个高斯分布,其中:
α
t
=
1
−
β
t
α
ˉ
t
=
Π
s
=
0
t
α
s
\alpha_t = 1 - \beta_t \\ \bar{\alpha}_t = \Pi_{s=0}^t \alpha_s
αt=1−βtαˉt=Πs=0tαs
当
α
ˉ
t
\bar{\alpha}_t
αˉt 取值趋近于0时,
z
t
z_t
zt 可以视为一个标准的高斯分布,在DiffusionDet中,
β
1
:
T
\beta_{1:T}
β1:T取了一系列零到一,且逐渐变大的值,以下是生成
β
\beta
β 的代码,这里我们取
T
=
1000
T=1000
T=1000,即共采样
1000
1000
1000 步
q
(
z
t
∣
z
t
−
1
)
:
=
N
(
z
t
;
1
−
β
t
z
t
−
1
,
β
t
I
)
q(z_t | z_{t-1}) := \mathcal{N} (z_{t}; \sqrt{1 - \beta_t} z_{t-1}, \beta_t \bf{I} )
q(zt∣zt−1):=N(zt;1−βtzt−1,βtI)
接下来计算
α
ˉ
t
\sqrt{\bar{\alpha}_{t}}
αˉt,
1
−
α
ˉ
t
\sqrt{1 - \bar{\alpha}_{t}}
1−αˉt,
log
(
1
−
α
ˉ
t
)
\log{(1-\bar{\alpha}_{t})}
log(1−αˉt),
1
α
ˉ
t
\frac{1}{\sqrt{\bar{\alpha}_{t}}}
αˉt1和
1
α
ˉ
t
−
1
\sqrt{\frac{1}{\bar{\alpha}_t} - 1}
αˉt1−1
DDPM文中假设,后验分布
q
(
z
t
−
1
∣
z
t
,
z
0
)
q(z_{t-1} | z_t, z_0)
q(zt−1∣zt,z0)也是高斯分布,有:
q
(
z
t
−
1
∣
z
t
,
z
0
)
=
N
(
z
t
−
1
;
μ
~
(
z
t
,
z
0
)
,
β
t
~
I
)
q(z_{t-1} | z_t, z_0) = \mathcal{N} (z_{t-1} ; \tilde{\mu}(z_t, z_0), \tilde{\beta_t} \bm{I})
q(zt−1∣zt,z0)=N(zt−1;μ~(zt,z0),βt~I)
算式整理后有:
μ
~
t
(
z
t
,
z
0
)
=
α
ˉ
t
−
1
β
t
1
−
α
ˉ
t
z
0
+
α
t
(
1
−
α
ˉ
t
−
1
)
1
−
α
ˉ
t
z
t
\tilde{\mu}_t(z_t, z_0) = \frac{ \sqrt{\bar{\alpha}_{t-1}} \beta_t }{ 1 - \bar{\alpha}_t } z_{0} + \frac { \sqrt{\alpha_t} (1 - \bar{\alpha}_{t-1}) } { 1 - \bar{\alpha}_t } z_{t}
μ~t(zt,z0)=1−αˉtαˉt−1βtz0+1−αˉtαt(1−αˉt−1)zt
β
~
t
=
1
−
α
ˉ
t
−
1
1
−
α
ˉ
t
β
t
\tilde{\beta}_{t} = \frac { 1 - \bar{\alpha}_{t-1} } { 1 - \bar{\alpha}_t } \beta_{t}
β~t=1−αˉt1−αˉt−1βt
接下来的几行代码用来计算这几个系数:
# calculations for posterior q(x_{t-1} | x_t, x_0)# equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t)
posterior_variance = betas *(1.- alphas_cumprod_prev)/(1.- alphas_cumprod)
self.register_buffer('posterior_variance', posterior_variance)# log calculation clipped because the posterior variance is 0 at# the beginning of the diffusion chain
self.register_buffer('posterior_log_variance_clipped',
torch.log(posterior_variance.clamp(min=1e-20)))
self.register_buffer('posterior_mean_coef1',
betas * torch.sqrt(alphas_cumprod_prev)/(1.- alphas_cumprod))
self.register_buffer('posterior_mean_coef2',(1.- alphas_cumprod_prev)* torch.sqrt(alphas)/(1.- alphas_cumprod))