从零点五开始的深度学习笔记——VAE(Variational AutoEncoder) （三）VAE的简单实现

学习笔记链接

从零点五开始的深度学习笔记——VAE(Variational AutoEncoder) （一） 预备知识
从零点五开始的深度学习笔记——VAE(Variational AutoEncoder) （二）概率角度理解VAE结构

1. 预备知识

1.1 关于采样

1.1.1 蒙特卡罗模拟

蒙特卡罗，蒙特卡洛，Monte Carlo是一个赌场的名字，这名字起得就很有概率统计学的意思。部分参考苏剑林. (Mar. 28, 2018). 《变分自编码器（二）：从贝叶斯观点出发 》[Blog post]. Retrieved from https://spaces.ac.cn/archives/5343 教程。下面我们将介绍如何求一个满足$p(x)$概率分布的随机变量$x$的期望的过程，首先在连续空间中分析，${\mathbb{E}}_{x\sim p(x)}\left[x\right]=\int xp(x)dx$, 如果使用计算机来硬算的话，一个简单的方式是将他离散化再累加起来：
$$\begin{array}{c}{\mathbb{E}}_{x\sim p(x)}\left[x\right]\approx \sum _i{x}_{i}p(x_i)(x_i-x_{i-1})\end{array}$$
根据统计学中的大数定理，一个随机变量的期望等于这个随机变量在$n\to \infty$次试验后，所有取值的平均值。所以计算机可以对采样进行模拟，然后将所有取值求平均，从而得到该变量的期望：
$$\begin{array}{c}{\mathbb{E}}_{x\sim p(x)}\left[x\right]\approx \frac{1}{n}\sum _{i=1}^n{x}_i,x_i\sim p(x)\end{array}$$
更一般的，
$$\begin{array}{c}{\mathbb{E}}_{x\sim p(x)}[f(x)]\approx \frac{1}{n}\sum _{i=1}^{n}f(x_i),x_i\sim p(x)\end{array}$$

1.1.2 重要性采样

上面解决了期望计算的问题，但还需要对$x$根据概率分布$p(x)$进行采样，这就不得不提重要性采样， MCMC采样等从一个已知分布中对随机变量进行采样的采样方法了。后续涉及到再补充。

1.2 VAE模型的假设

实际上上一篇博文从零点五开始的深度学习笔记——VAE(Variational AutoEncoder) （二）概率角度理解VAE结构的目标函数是一般化的，并非VAE特有，根据理论公式的指导，VAE对网络的具体构造和实现进行了一定的假设，使之实现网络生成的功能。回顾一下这条理想状态下的VAE的损失函数：
$$\begin{array}{c}{\mathbb{D}}_{KL}\left(p(x,z)\left|\right|q(x,z)\right)={\mathbb{E}}_{x\sim p(x)}\left[{\mathbb{E}}_{z\sim p(z|x)}\right[-log(q(x|z))]+{\mathbb{D}}_{KL}\left(p(z|x)||q(z)\right)]+const\end{array}$$
为什么说这条公式具有指导意义？因为他告诉了我们所有相关的随机变量是在什么概率分布下采样出来的（虽然采样的概率分布可能是未知的），且如果设计了一个网络，这个网络应该需要逼近哪些项，如果模型极限效果提不上去，是由于我们做了什么理想的假设使得模型于这纷繁复杂的世界之间存在reality gap。当然，这些假设的存在往往是为了简化问题的解决难度，如下文所示：

1.2.1 关于采样

上面我们简单提到了蒙特卡洛采样的思想。如果我们数据足够庞大的话，上面的损失函数可以简写为：
$$\begin{array}{c}{\mathbb{D}}_{KL}\left(p(x,z)\left|\right|q(x,z)\right)={\mathbb{E}}_{z\sim p(z|x)}[-log(q(x|z))]+{\mathbb{D}}_{KL}\left(p(z|x)||q(z)\right)\end{array}$$
为什么？因为我们的训练输入数据就是已经遵循了某个未知的规律采回来的，狗有狗的样子，采样回来不会变成猫的样子。只要不是人工合成的数据，我们所获得的训练集在冥冥之中已经遵循了某种规律。也因此，我们可以把这个理想的损失函数最外层的剥离。这一步假设成立的前提是训练集足够庞大，能够比较好地体现采样分布。

继续观察这个函数，涉及到了几个部分，理想模型中的$p(z|x)$, 待估计网络的$q(z)$和$q(x|z)$。不难理解，这三项正好对应了，编码器，隐空间(latent space, latent vector, bottle neck，都是它)，以及解码器。

1.2.2 编码器$p(z|x)$部分

简化后的损失函数还是具有一定的缺陷，涉及到蛋鸡问题。为了拟合这个未知模型需要已知理想的$p(z|x)$。我猜研究人员把公式推到这里后，实在是不能忍，想要窥探VAE的真容却在此刻吃了只苍蝇。所以只能硬着头皮把这一项用万物皆可神经网络拟合来继续向前探索了。
$p(z|x)$编码器也使用神经网络拟合。由于最终需要拟合的是一个概率分布用于后续的采样，从而得到隐变量$z$，所以继续引入了假设，假设这个概率分布是一个多变量高斯函数，即多变量正态分布函数。则，神经网络的输出为隐变量$z$各个维度的均值和方差。总而言之，就是每一个维度的$z$都是一个高斯函数, 即$z_i\sim \mathcal{N}({\mu }_i(x),{\sigma }_i^2(x))$
$$\begin{array}{c}\begin{array}{rl}p(z|x)=& \mathcal{N}(\mu (x),diag({\sigma }_i^2(x)\left)\right),i=1,2,...,k,z\in {\mathbb{R}}^k\\ =& \frac{1}{{\prod }_{i=1}^k\sqrt{2\pi {\sigma }_i^2(x)}}exp\left(-\frac{1}{2}\sum _{i=1}^k\frac{(z_i-{\mu }_i(x){)}^2}{{\sigma }_i^2(x)}\right)\end{array}\end{array}$$

1.2.3 隐变量$q(z)$部分

$q(z)$这一部分就是完全可控的了，可以由我们自己设计。那当然是一切从简，所以，直接假定这个分布是标准正态分布，可以让整个世界都变得很美好。

$$p(z)=\mathcal{N}\left(0,I\right)$$

回顾从零点五开始的深度学习笔记——VAE(Variational AutoEncoder) （一） 预备知识 中我们推导得到的公式：
$$\begin{array}{c}{\mathbb{D}}_{KL}(P||Q)=\frac{1}{2}\left[log\frac{|{\Sigma }_2|}{|{\Sigma }_1|}-k+tr({\Sigma }_2^{-1}{\Sigma }_1)+({\mu }_2-{\mu }_1{)}^T{\Sigma }_2^{-1}({\mu }_2-{\mu }_1)\right]\end{array}$$
将$q(z)$和$p(z|x)$代入到上式中，可得：

$${\mathbb{D}}_{KL}(p(z|x)\left|\right|q(z))=\frac{1}{2}\sum _{i=1}^k({\sigma }_i^2(x)+{\mu }_i^2(x)-log({\sigma }_i^2(x))-1)$$

1.2.4 解码器$q(x|z)$部分

这部分的理解参考了苏剑林. (Mar. 28, 2018). 《变分自编码器（二）：从贝叶斯观点出发 》[Blog post]. Retrieved from https://spaces.ac.cn/archives/5343的分析。解码器的模型候选方案需要输出的是一个容易计算，且满足概率分布函数积分为1的归一化约束。因此，可选择的方案不多。有伯努利分布和正态分布模型。

• 正态分布模型
  和$p(z|x)$相同，
  $$q(x|z)=\mathcal{N}(\bar{\mu }(z),diag({\bar{\sigma }}_i^2(z))),i=1,2,...,n,x\in {\mathbb{R}}^n$$
  则，损失函数的第一项可展开为：
  $$-log(q(x|z))=\frac{1}{2}\sum _{i=1}^n\frac{(z_i-{\bar{\mu }}_i(z){)}^2}{{\bar{\sigma }}_i^2(z)}+\frac{1}{2}\sum _{i=1}^{n}log({\bar{\sigma }}_i^2(z))+\frac{n}{2}log(2\pi )$$

1.2.5 小结

因此，按照这何种结构选择，最终的损失函数的计算可简化为：
$$\begin{array}{c}\begin{array}{rl}& {\mathbb{D}}_{KL}(p(x,z)||q(x,z))\\ =& {\mathbb{E}}_{z\sim p(z|x)}[\frac{1}{2}\sum _{i=1}^n\frac{(z_i-{\bar{\mu }}_i(z){)}^2}{{\bar{\sigma }}_i^2(z)}+\frac{1}{2}\sum _{i=1}^{n}log({\bar{\sigma }}_i^2(z))+\frac{n}{2}log(2\pi )]\\ & +\frac{1}{2}\sum _{i=1}^k({\sigma }_i^2(x)+{\mu }_i^2(x)-log({\sigma }_i^2(x))-1)\end{array}\end{array}$$
注意，这里的算是函数跟一般的损失函数不太一样，有个期望在这里，在训练模型的时候，需要对这一项进行采样，从实践结果来看，可以用采样一次的结果来表示期望的值。因此, 参考苏剑林. (Mar. 28, 2018). 《变分自编码器（二）：从贝叶斯观点出发 》[Blog post]. Retrieved from https://spaces.ac.cn/archives/5343，整个VAE的损失函数的一般性写法是：
$$\begin{array}{c}\begin{array}{rl}\mathcal{L}=& {\mathbb{D}}_{KL}(p(x,z)||q(x,z))\\ =& {\mathbb{E}}_{z\sim p(z|x)}[-log(q(x|z))+{\mathbb{D}}_{KL}(p(z|x)\left|\right|q(z)\left)\right],z\sim p(z|x)\end{array}\end{array}$$

2. VAE的实现

2.1 重参数化

重参数化 (reparameterizatin) 是VAE编程实现很重要的小技巧，有了它才能够让网络反向传播，计算梯度，更新网络权值。它使用到了正态分布采样的一个等价的方法：在标准正态分布$\mathcal{N}(0,\mathbf{I})$中采样得到$ϵ$后，对采样结果进行放缩$\sigma$和位移$\mu$可以使得到的采样结果($\mu +\sigma ϵ$)与另一个相关的正态分布$\mathcal{N}(\mu ,{\sigma }^2)$采样结果一致。这种采样方式之所以称为小技巧是因为这种操作是为了适配当前编程软件自动梯度求解功能所做的操作。

2.2 以MNIST手写数字图片为例

2.2.1 MNIST数据下载

'''
Author       : Dianye Huang
Date         : 2022-08-23 10:04:45
LastEditTime: 2022-08-26 22:02:34
Description  : 
'''

from torch.utils.data import DataLoader

import torchvision 
from torchvision.datasets import mnist
import torchvision.transforms as transforms

class ExpDataLoader(object):
    def __init__(self) -> None:
        self.to_pil_image = transforms.ToPILImage()

    def vis_img(self, img):
        vis = self.to_pil_image(img)
        vis.show()
        
    def vis_grid_imgs(self, imgs, nrow=8):
        grid = torchvision.utils.make_grid(imgs, nrow=nrow)
        self.vis_img(grid)

    def get_mnist_dataset(self, dir='./data', ):
        train_set=mnist.MNIST(dir, train=True, 
                                transform=torchvision.transforms.ToTensor(), 
                                download=True)
        test_set=mnist.MNIST(dir, train=False, 
                                transform=torchvision.transforms.ToTensor(), 
                                download=True)
        return train_set, test_set
    
    def get_mnist_dataloader(self, dir='./data', batch_size = 16):
        mnist_train_ds, mnist_test_ds = self.get_mnist_dataset(dir=dir)
        train_loader = DataLoader(dataset=mnist_train_ds, 
                                    batch_size=batch_size, 
                                    shuffle=True)
        test_loader = DataLoader(dataset=mnist_test_ds, 
                                    batch_size=batch_size, 
                                    shuffle=False)
        return train_loader, test_loader

下载数据生成dataloader

if __name__ == '__main__':
    # load data
    exp_dataloader = ExpDataLoader()
    data_dir = '/home/dianye/DNN_ws/CSDN_tutorials/VAEs'
    train_loader, test_loader = exp_dataloader.get_mnist_dataloader(dir=data_dir, batch_size=128)

2.2.2 传统的的自编码器

从 https://avandekleut.github.io/vae/ 中复制过来的一段最最原始的自编码器代码，简单粗暴地拟合输入输出。

import torch
from torch import nn
import torch.nn.functional as F

'''
Typical Auto Encoder:
https://avandekleut.github.io/vae/
'''
class Encoder(nn.Module):
    def __init__(self, latent_dims):
        super(Encoder, self).__init__()
        self.linear1 = nn.Linear(784, 512)
        self.linear2 = nn.Linear(512, latent_dims)

    def forward(self, x):
        x = torch.flatten(x, start_dim=1)
        x = F.relu(self.linear1(x))
        return self.linear2(x)
    
class Decoder(nn.Module):
    def __init__(self, latent_dims):
        super(Decoder, self).__init__()
        self.linear1 = nn.Linear(latent_dims, 512)
        self.linear2 = nn.Linear(512, 784)

    def forward(self, z):
        z = F.relu(self.linear1(z))
        z = torch.sigmoid(self.linear2(z))
        return z.reshape((-1, 1, 28, 28))

class Autoencoder(nn.Module):
    def __init__(self, latent_dims):
        super(Autoencoder, self).__init__()
        self.encoder = Encoder(latent_dims)
        self.decoder = Decoder(latent_dims)

    def forward(self, x):
        z = self.encoder(x)
        return self.decoder(z)

开始训练，下面的代码主要参考了https://avandekleut.github.io/vae/，加了点小修改。

'''
Author       : Dianye Huang
Date         : 2022-08-23 10:04:45
LastEditTime: 2022-08-26 22:08:16
Description  : 
'''

import torch
import torchvision
from vae_utils import ExpDataLoader
from vae_zoo import Autoencoder

import matplotlib.pyplot as plt
import numpy as np

# device = 'cuda' if torch.cuda.is_available() else 'cpu'
from tqdm import tqdm
device = 'cpu'
def train_ae(autoencoder, data, epochs=20):
    opt = torch.optim.Adam(autoencoder.parameters())
    for epoch in range(epochs):
        print('Epoch:', epoch)
        for x, y in tqdm(data):
            x = x.to(device) # GPU
            opt.zero_grad()
            x_hat = autoencoder(x)
            loss = ((x - x_hat)**2).sum()
            loss.backward()
            opt.step()
    return autoencoder

def plot_latent(autoencoder, data, num_batches=100):
    for i, (x, y) in enumerate(data):
        z = autoencoder.encoder(x.to(device))
        z = z.to('cpu').detach().numpy()
        plt.scatter(z[:, 0], z[:, 1], c=y, cmap='tab10')
        if i > num_batches:
            plt.colorbar()
            break

def plot_reconstructed(autoencoder, r0=(-5, 10), r1=(-10, 5), n=12):
    w = 28
    img = np.zeros((n*w, n*w))
    for i, y in enumerate(np.linspace(*r1, n)):
        for j, x in enumerate(np.linspace(*r0, n)):
            z = torch.Tensor([[x, y]]).to(device)
            x_hat = autoencoder.decoder(z)
            x_hat = x_hat.reshape(28, 28).to('cpu').detach().numpy()
            img[(n-1-i)*w:(n-1-i+1)*w, j*w:(j+1)*w] = x_hat
    plt.imshow(img, extent=[*r0, *r1])


if __name__ == '__main__':
    # load data
    exp_dataloader = ExpDataLoader()
    data_dir = '/home/dianye/DNN_ws/CSDN_tutorials/VAEs'
    train_loader, test_loader = exp_dataloader.get_mnist_dataloader(dir=data_dir, batch_size=128)
	
	# start training
	device = 'cpu'
	latent_dims = 2
	autoencoder = Autoencoder(latent_dims).to(device) # GPU
	autoencoder = train_ae(autoencoder, train_loader, epochs=20)
	
	# visualize result
	plt.figure(1)
	plot_latent(autoencoder, train_loader)
	plt.figure(2)
	plot_reconstructed(autoencoder)
	plt.pause(0)

2.2.4 训练结果

以下是训练了20个epoch之后得到的结果

2.3 VAE

2.3.1 模型

VAE编程的时候，需要损失函数的写法。第一点是，bottleneck部分，两个全连接层输出的分别是均值$\mu$和方差$\log ({\sigma }^2)$。因此在计算KL散度的时候会有exp的运算。对于模型拟合误差项为什么是一个输入输出的平方差之和，可以观察我们损失函数的第一项：
$$\begin{array}{c}\begin{array}{rl}& {\mathbb{D}}_{KL}(p(x,z)||q(x,z))\\ =& {\mathbb{E}}_{z\sim p(z|x)}[\frac{1}{2}\sum _{i=1}^n\frac{(z_i-{\bar{\mu }}_i(z){)}^2}{{\bar{\sigma }}_i^2(z)}+\frac{1}{2}\sum _{i=1}^{n}log({\bar{\sigma }}_i^2(z))+\frac{n}{2}log(2\pi )]\\ & +\frac{1}{2}\sum _{i=1}^k({\sigma }_i^2(x)+{\mu }_i^2(x)-log({\sigma }_i^2(x))-1)\end{array}\end{array}$$
$p(z|x)$是解码器，采样服从的分布，而我们实际计算前向传播的时候，并没有用到方差，而是认为方差是个常数，直接以概率1将$z$不做任何处理直接扔到解码器里。这样，优化损失函数的第一项，就相当于优化$\frac{1}{2}{\sum }_{i=1}^n\frac{(z_i-{\bar{\mu }}_i(z){)}^2}{{\bar{\sigma }}_i^2(z)}$中的平方xiang，即$\frac{1}{2{\bar{\sigma }}_i^2(z)}{\sum }_{i=1}^n(z_i-{\bar{\mu }}_i(z){)}^2=const\times \Vert z-decoder(z){\Vert }^2$。

'''
Author       : Dianye Huang
Date         : 2022-08-23 10:04:45
LastEditTime : 2022-08-27 01:34:39
Description  : 
'''

import torch
from torch import nn
import torch.nn.functional as F

'''
Typical Variational Auto Encoder
'''
class VanillaVAE(nn.Module):
    def __init__(self,
                in_channels:int = 784, # 28*28
                latent_dim: int = 2,
                hidden_dims: list = [512]
                ) -> None:
        super(VanillaVAE, self).__init__()
        
        self.in_channels = in_channels
        
        # Build Encoder
        modules = []
        for h_dim in hidden_dims:
            modules.append(
                nn.Sequential(
                    nn.Linear(in_channels, h_dim),
                    nn.ReLU()
                )
            )
            in_channels = h_dim
        self.encoder = nn.Sequential(*modules)

        # Bottle Neck
        self.latent_dim = latent_dim
        self.fc_mu = nn.Linear(hidden_dims[-1], latent_dim)
        self.fc_var = nn.Linear(hidden_dims[-1], latent_dim)
        
        # Build Decoder
        modules = []
        in_ch = latent_dim
        hidden_dims.reverse()
        for h_dim in hidden_dims:
            modules.append(
                nn.Sequential(
                    nn.Linear(in_ch, h_dim),
                    nn.ReLU()
                )
            )
            in_ch = h_dim
        self.decoder = nn.Sequential(*modules)

        # Output Layer
        self.output_layer = nn.Sequential(
                            nn.Linear(hidden_dims[-1], self.in_channels),
                            nn.Sigmoid())
    
    def encode(self, input:torch.tensor):
        return self.encoder(input)

    def bottleneck(self, input:torch.tensor):
        mu = self.fc_mu(input)
        log_var = self.fc_var(input)
        return self.reparameterize(mu, log_var)
    
    def decode(self, z: torch.tensor):
        return self.decoder(z)
    
    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std) # 返回一个和输入大小相同的张量，其由均值为0、方差为1的标准正态分布填充
        return eps*std + mu
        
    def forward(self, input: torch.tensor):
        x = torch.flatten(input, start_dim=1)
        out = self.encoder(x)
        mu = self.fc_mu(out)
        logvar = self.fc_var(out)
        z = self.reparameterize(mu, logvar)
        out = self.decoder(z)
        out = self.output_layer(out)
        return out, x, mu, logvar
    
    def loss_function(self, x_hat, x, mu, logvar):
        D_KL = 0.5*(torch.exp(logvar) + mu**2 - logvar - 1).sum()
        recon_loss = 10*((x_hat - x)**2).sum()
        return recon_loss + D_KL 

2.3.2 训练

'''
Author       : Dianye Huang
Date         : 2022-08-23 10:04:45
LastEditTime: 2022-08-27 00:43:14
Description  : 
'''

import torch
from vae_utils import ExpDataLoader
from vae_zoo import Autoencoder, VanillaVAE

import matplotlib.pyplot as plt
import numpy as np
from tqdm import tqdm

device = 'cpu'
if __name__ == '__main__':
    # load data
    exp_dataloader = ExpDataLoader()
    data_dir = '/home/dianye/DNN_ws/CSDN_tutorials/VAEs'
    train_loader, test_loader = exp_dataloader.get_mnist_dataloader(dir=data_dir, batch_size=128)
    
    # start training variational auto encoder
    vae = VanillaVAE(latent_dim=2).to(device)
    opt = torch.optim.Adam(vae.parameters())
    for epoch in range(20):
        pbar = tqdm(train_loader, desc='description')
        for x, y in pbar:
            x = x.to(device) 
            opt.zero_grad()
            x_hat, x, mu, logvar = vae(x)
            loss = vae.loss_function(x_hat, x, mu, logvar)
            loss.backward()
            opt.step()
            pbar.set_description(f"Epoch: {epoch+1}, loss: {round(float(loss.to('cpu').detach().numpy()),3)}")
    

    plt.figure(1)
    with torch.no_grad():
        for i, (x, y) in enumerate(train_loader):
            z_tmp = vae.encoder(torch.flatten(x.to(device), start_dim=1))
            z = vae.bottleneck(z_tmp)
            z = z.to('cpu').detach().numpy()
            plt.scatter(z[:, 0], z[:, 1], c=y, cmap='tab10')
            if i > 100:
                plt.colorbar()
                break
    
    plt.figure(2)
    with torch.no_grad():
        r0=(-5, 10)
        r1=(-10, 5)
        n=12
        w = 28
        img = np.zeros((n*w, n*w))
        for i, y in enumerate(np.linspace(*r1, n)):
            for j, x in enumerate(np.linspace(*r0, n)):
                z = torch.Tensor([[x, y]]).to(device)
                x_tmp = vae.decoder(z)
                x_hat = vae.output_layer(x_tmp)
                x_hat = x_hat.reshape(28, 28).to('cpu').detach().numpy()
                img[(n-1-i)*w:(n-1-i+1)*w, j*w:(j+1)*w] = x_hat
        plt.imshow(img, extent=[*r0, *r1])
        
    plt.show()

2.3.3 训练结果

以上。VAE部分理解得差不多就行了，后续的博客将会介绍conditional VAE。

祝周末愉快！

2022年8月27日
Dianye Huang