VAE 和 VQ-VAE

1
# VAE 伪代码
2
1  x = input_image
3
2  mu, logvar = Encoder(x)
4
3  z = mu + exp(0.5 * logvar) * eps # eps ~ N(0, I)
5
4  x_hat = Decoder(z)
6

7
5  loss_recon = reconstruction_loss(x_hat, x)
8
6  loss_kl = KL(q(z|x) || N(0, I))
9

10
# 然后 VQ-VAE：
11
1  x = input_image
12
2  z_e = Encoder(x)
13
3  z_q = argmin_{z in codebook} distance(z_e, z) # 在 codebook 里找离 z_e 最近的向量
14
4  x_hat = Decoder(z_q)
15

4 collapsed lines
16
5  loss_recon = reconstruction_loss(x_hat, x)
17
6  loss_codebook = || stop_grad(z_e) - z_q ||^2
18
7  loss_commit = || z_e - stop_grad(z_q) ||^2
19
8  loss = loss_recon + loss_codebook + beta * loss_commit