变分自编码器VAE实现MNIST数据集生成by Pytorch

参考：
原文：Auto-Encoding Variational Bayes

Recap自编码器：

自编码器中，需要输入一个原始图片，原始图片经过编码之后得到一个隐向量，隐向量解码产生原图片对应的图片。在这种情况下，只能生成原图片对应的图片而无法任意生成新的图片，因为隐向量都是原始图片确定的。

变分自编码器VAE

引入变分自编码器（Variational autoencoder）可以在遵循某一分布下随机产生一些隐向量来生成与原始图片不相同的图片，而不需要预先给定原始图片。为了达到这个目的，需要在编码过程增加限制，使得生成的隐向量能够粗略地遵循标准正态分布。
实际情况下，需要在模型的准确率与隐向量服从标准正态分布之间做一个权衡。模型的准确率就是指解码器生成的图片与原图片的相似程度；隐向量分布采用KL散度来衡量与标准正态分布之间的误差。两部分误差之和作为总体的误差来优化。

这里VAE使用了重参数化这个技巧来KL散度的计算。编码器不再是生成一个隐向量，而是生成正态分布的均值和标准差（若是多维正态分布，会有多个均值和标准差），然后根据这两个统计量下的分布抽样生成隐含向量。因为我们想要使得隐含向量服从标准正态分布，即均值为0，标准差为1，通过优化KL散度来使得分布逼近标准正态分布。

同理，解码器阶段，根据给定的隐变量$z$来生成多元正态分布的均值${\mu }_{x1},{\mu }_{x2}$标准差${\sigma }_{z1},{\sigma }_{z2}$,根据该分布抽样生成数值$x_1,x_2$

将编码器和解码器综合在一起：

设编码器的概率分布为$q_{\varphi }(z|x)$,解码器的概率分布为$p_{\theta }(x|z)$

误差推导：

这里需要用到正态分布之间的KL散度，直接给出公式，推导见参考文献：
单元正态分布：
$$KL({\mu }_1,{\mu }_2,{\sigma }_1,{\sigma }_2)=\log \frac{{\sigma }_2}{{\sigma }_1}+\frac{{\sigma }_1^2+({\mu }_1-{\mu }_2{)}^2}{2{\sigma }^2}-\frac{1}{2}$$
多$n$元正态分布：
$$KL({\mathbf{\mu }}_1,{\mathbf{\mu }}_2,{\Sigma }_1,{\Sigma }_2)=\frac{1}{2}[\log \frac{\det {\Sigma }_2}{\det {\Sigma }_1}-n+tr({\Sigma }_2^{-1}{\Sigma }_1)+({\mathbf{\mu }}_2-{\mathbf{\mu }}_1{)}^T{\Sigma }_2^{-1}({\mathbf{\mu }}_2-{\mathbf{\mu }}_1)]$$
由此可得到：
$$-D_{KL}(q(z|x^i)||p(z))=\frac{1}{2}\sum _{j=1}^J(1+\log (({\sigma }_{zj}^{(i)}{)}^2)-({\mu }_{zj}^{(i)}{)}^2-({\sigma }_{zj}^{(i)}{)}^2)$$
通过从分布$q(z|x^{(i)})$抽样来近似${\mathbb{E}}_{q(z|x^{(i)})}$。抽样$L$次，得到$z^{(i,l)},l=1,2,...,L$,$L$通常非常小，通常取1
$${\mathbb{E}}_{q(z|x^{(i)})}(\log (p(x^{(i)}|z)))=\frac{1}{L}\sum _{l=1}^L\log p(x^{(i)}|z^{(i,l)})=\frac{1}{L}\sum _{l=1}^L\sum _{j=1}^D\frac{1}{2}\log {\sigma }_{xj}^2+\frac{(x_j^i-{\mu }_{xj})}{2{\sigma }_{xj}^2}$$
其中$D$代表样本$x^{(i)}$的维度,每个数$x_j^{(i)}$都对应一个正态分布$\mathcal{N}({\mu }_{xj},{\sigma }_{xj}^2)$。

Pytorch实现MNIST数据集生成

在本实例中，生成器最后的输出不是均值和方差，而是图片向量。所以重构误差看做为生成图片和原始图片的误差。在这里使用的是binary cross entropy，即BCE误差，因为图片中的值都是(0,1)。当然也可以使用平方误差。
$$BCE=-\sum _{i=1}^n\sum _{j=1}^d[y_j^{(i)}\log x_j^{(i)}+(1-y_j^{(i)})\log (1-x_j^{(i)})]$$

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import torchvision.transforms as transforms
from torchvision.utils import save_image


def loss_function(recon_x, x, mu, logvar):
    """
    :param recon_x: generated image
    :param x: original image
    :param mu: latent mean of z
    :param logvar: latent log variance of z
    """
    BCE_loss = nn.BCELoss(reduction='sum')
    reconstruction_loss = BCE_loss(recon_x, x)
    KL_divergence = -0.5 * torch.sum(1+logvar-torch.exp(logvar)-mu**2)
    #KLD_ele = mu.pow(2).add_(logvar.exp()).mul_(-1).add_(1).add_(logvar)
    #KLD = torch.sum(KLD_ele).mul_(-0.5)
    print(reconstruction_loss, KL_divergence)

    return reconstruction_loss + KL_divergence


class VAE(nn.Module):
    def __init__(self):
        super(VAE, self).__init__()
        self.fc1 = nn.Linear(784, 400)
        self.fc2_mean = nn.Linear(400, 20)
        self.fc2_logvar = nn.Linear(400, 20)
        self.fc3 = nn.Linear(20, 400)
        self.fc4 = nn.Linear(400, 784)

    def encode(self, x):
        h1 = F.relu(self.fc1(x))
        return self.fc2_mean(h1), self.fc2_logvar(h1)

    def reparametrization(self, mu, logvar):
        # sigma = 0.5*exp(log(sigma^2))= 0.5*exp(log(var))
        std = 0.5 * torch.exp(logvar)
        # N(mu, std^2) = N(0, 1) * std + mu
        z = torch.randn(std.size()) * std + mu
        return z

    def decode(self, z):
        h3 = F.relu(self.fc3(z))
        return torch.sigmoid(self.fc4(h3))

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparametrization(mu, logvar)
        return self.decode(z), mu, logvar


transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5]),
])

trainset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True)

testset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False)

vae = VAE()
optimizer = torch.optim.Adam(vae.parameters(), lr=0.0003)

# Training
def train(epoch):
    vae.train()
    all_loss = 0.
    for batch_idx, (inputs, targets) in enumerate(trainloader):
        inputs, targets = inputs.to('cpu'), targets.to('cpu')
        real_imgs = torch.flatten(inputs, start_dim=1)

        # Train Discriminator
        gen_imgs, mu, logvar = vae(real_imgs)
        loss = loss_function(gen_imgs, real_imgs, mu, logvar)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        all_loss += loss.item()
        print('Epoch {}, loss: {:.6f}'.format(epoch, all_loss/(batch_idx+1)))
        # Save generated images for every epoch
    fake_images = gen_imgs.view(-1, 1, 28, 28)
    save_image(fake_images, 'MNIST_FAKE/fake_images-{}.png'.format(epoch + 1))



for epoch in range(20):
    train(epoch)

torch.save(vae.state_dict(), './vae.pth')

运行上述代码20轮所产生的的图片：

VAE和自编码器有一样的缺点，根据均平方误差计算的图片会比较模糊，之后出现的对抗生成网络则解决了这个问题。