变分自编码器+要点综述+代码实现+生成图片

1.VAE的结构

变分自编码器（Variational Autoencoders）是由Diederik Kingma和Max Welling在2014年提出来的。

1.1 网络结构

VAE的基本结构如下图所示,来自《Hands On ML》: Figure 15-11:Variational autoencoder (left), and an instance going through it (right)。
上述的Hidden1,Hidden2可以使密集层或者卷基层。

模型结构上来看Variational Autoencoders 和Autoencoders 的区别主要在Coding Layer：VAE希望在模型输入和输出尽量相同的时候，同时编码层拟合一个正态分布（例如标准正态分布）。VAE希望编码层最终的输出向量仿佛是从一个正态分布抽样出来的，同时还和输入很像。

其实现方法是VAE编码层拟合了一个正态分布的均值向量μ和方差向量（log(σ^2)），然后用二者计算（这个计算过程又叫reparameterize）得到一个带有随机性的向量。再对这个向量进行解码，希望解码之后的Outputs尽量和Inputs相同。

1.2 作为生成模型的VAE

当训练完一个VAE之后，随机生成若干正态分布向量（根据训练时候Coding Layer预设的正态分布函数，例如标准正态分布）。用训练好的VAE的解码部分进行计算。就可以得到若干“很像”原始数据的数据或者图片。

1.3 更多思想方面的细节

关于VAE的更多细节，请参考一篇写的很好的知文 变分自编码器VAE：原来是这么一回事 | 附开源代码 。

2.VAE的损失函数

【说明：VAE在代码实现的过程中，个人感觉最难理解的就是损失函数，其他部分和常规的神经网络差不多】

作为自编码器，损失函数肯定要考虑Inputs和Outputs的相似性，所以需要有Reconstruction Loss，这个比较好理解和实现——就是衡量输入和输出的差异。

变分自编码器由于其Coding Layer的特殊性，其在Coding Layer需要“拟合”一个标准正态分布。所以需要衡量这个“拟合的程度”。所以VAE的损失函数还需要一个Latent Loss。

VAE最终的损失函数 Loss=Reconstruction_Loss+Latent_Loss

2.1 重构损失-Reconstruction Loss

衡量输出和输出之间误差的方法应该很多。参考了几个资料，基本都是交叉熵来作为损失函数（直接用输出输出的均方误差似乎收敛的很慢）。
用tensorflow代码表示重构误差如下（其中X是样本输入，logits是VAE模型输出）：

reconstruction_loss = tf.reduce_sum(
			tf.nn.sigmoid_cross_entropy_with_logits(labels=X, logits=logits))

2.2 KL-Latent Loss

VAE的Coding Layer希望拟合一个标准正态分布，因此衡量这种“拟合的程度”最好理解的一种方法就是KL散度——KL散度给出两个分布的差异的度量。 KL散度越大，两个分布差别越大。

KL散度损失函数的推导公式，请参考变分自编码器KL散度的Latent Loss推导 。

2.3 ELBO-Latent Loss

【说明：下述的公式中，x可以理解为输入VAE的样本；z为隐变量，可以理解为VAE的encoder的输出；q(z)表示隐变量的目标分布；p(z|x)表示拟合出的分布；p(x)为样本数据的实际分布】

ELBO是指evidence lower bound。

2.4 更多理论细节

请参考 Notes on Variational Autoencoders 。

3. 代码实现

3.1 CVAE训练和生产图片-全部代码

下述代码实现，主要来自于 Tensorflow-Convolutional Variational Autoencoder的教程。教程中用的ELBO latent loss。自己在实验过程中，添加了KL Latent Loss的函数，看上去实验结果二者差距不是很大。

这个代码比较方便的地方是，如果需要把卷积层修改成其他类型的层（例如密集层），可以在类里面直接修改，很方便。

import tensorflow as tf
import time
import numpy as np
import matplotlib.pyplot as plt
from IPython import display

#################################################################
#################################################################
class CVAE(tf.keras.Model):
    def __init__(self, latent_dim):
        super(CVAE, self).__init__()
        self.latent_dim = latent_dim
        self.inference_net = tf.keras.Sequential(
            [
                tf.keras.layers.InputLayer(input_shape=(28, 28, 1)),
                tf.keras.layers.Conv2D(
                    filters=32, kernel_size=3, strides=(2, 2), activation='relu'),
                tf.keras.layers.Conv2D(
                    filters=64, kernel_size=3, strides=(2, 2), activation='relu'),
                tf.keras.layers.Flatten(),
                # No activation
                tf.keras.layers.Dense(latent_dim + latent_dim),
            ]
        )
        self.generative_net = tf.keras.Sequential(
            [
                tf.keras.layers.InputLayer(input_shape=(None,latent_dim)),
                tf.keras.layers.Dense(units=7 * 7 * 32, activation=tf.nn.relu),
                tf.keras.layers.Reshape(target_shape=(7, 7, 32)),
                tf.keras.layers.Conv2DTranspose(
                    filters=64,
                    kernel_size=3,
                    strides=(2, 2),
                    padding="SAME",
                    activation='relu'),
                tf.keras.layers.Conv2DTranspose(
                    filters=32,
                    kernel_size=3,
                    strides=(2, 2),
                    padding="SAME",
                    activation='relu'),
                # No activation
                tf.keras.layers.Conv2DTranspose(
                    filters=1, kernel_size=3, strides=(1, 1), padding="SAME"),
            ]
        )

    @tf.function
    def sample(self, eps=None):
        if eps is None:
          eps = tf.random.normal(shape=(100, self.latent_dim))
        return(self.decode(eps, apply_sigmoid=True))

    def encode(self, x):
        mean, logvar = tf.split(self.inference_net(x), num_or_size_splits=2, axis=1)
        return(mean, logvar)

    def reparameterize(self, mean, logvar):
        eps = tf.random.normal(shape=mean.shape)
        return(eps * tf.exp(logvar *0.5) + mean)

    def decode(self, z, apply_sigmoid=False):
        logits = self.generative_net(z)
        if apply_sigmoid:
          probs = tf.sigmoid(logits)
          return(probs)
        return(logits)

#################################################################
#### KL-Latent Loss
@tf.function
def vae_kl_loss(model,x):
    ## latent loss
    coder_mean,coder_gamma=model.encode(x)
    coder_sigma=tf.exp(0.5*coder_gamma)
    noise=tf.random.normal(shape=coder_gamma.shape)
    z=coder_mean+noise*coder_sigma
    latent_loss = 0.5 * tf.reduce_sum(
                        tf.exp(coder_gamma) + tf.square(coder_mean) - 1 - coder_gamma)
    ## reconstruction loss
    logits=model.decode(z)
    reconstruction_loss=tf.reduce_sum(
    	                tf.nn.sigmoid_cross_entropy_with_logits(logits=logits,labels=x))
    ##
    vae_loss_ = reconstruction_loss + latent_loss
    return(vae_loss_)

#### ELBO-Latent Loss
def log_normal_pdf(sample, mean, logvar, raxis=1):
    log2pi = tf.math.log(2.0*np.pi)
    return(tf.reduce_sum(
      -0.5*((sample-mean)**2.0*tf.exp(-logvar)+logvar+log2pi),
      axis=raxis))

@tf.function
def vae_elbo_loss(model, x):
    mean,logvar=model.encode(x)
    z=model.reparameterize(mean,logvar)
    x_logit=model.decode(z)
    #
    cross_ent=tf.nn.sigmoid_cross_entropy_with_logits(logits=x_logit,labels=x)
    logpx_z=-tf.reduce_sum(cross_ent,axis=[1, 2, 3])
    logpz=log_normal_pdf(z,0.0,0.0)
    logqz_x=log_normal_pdf(z,mean,logvar)
    return(-tf.reduce_mean(logpx_z+logpz-logqz_x))

#### update the model with gradients
@tf.function
def compute_apply_gradients(model,x,optimizer):
    with tf.GradientTape() as tape:
        loss=vae_kl_loss(model,x)
    gradients=tape.gradient(loss,model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

#################################################################
#################################################################
#### load mnist data
(train_images, _), (test_images, _) = tf.keras.datasets.mnist.load_data()
#### data process
train_images = train_images.reshape(train_images.shape[0], 28, 28, 1).astype('float32')
test_images = test_images.reshape(test_images.shape[0], 28, 28, 1).astype('float32')
## Normalizing the images to the range of [0., 1.]
train_images /= 255.0
test_images /= 255.0
## Binarization
train_images[train_images >= 0.5] = 1.0
train_images[train_images < 0.5] = 0.0
test_images[test_images >= 0.5] = 1.0
test_images[test_images < 0.5] = 0.0
##
TRAIN_BUF = 60000
BATCH_SIZE = 100
TEST_BUF = 10000
##
train_dataset = tf.data.Dataset.from_tensor_slices(train_images).shuffle(TRAIN_BUF).batch(BATCH_SIZE)
test_dataset = tf.data.Dataset.from_tensor_slices(test_images).shuffle(TEST_BUF).batch(BATCH_SIZE)
####  训练参数
epochs = 50
latent_dim = 50
num_examples_to_generate = 16

## 生成图片的隐变量，用标准正态分布随机生成
random_vector_for_generation = tf.random.normal(
    shape=[num_examples_to_generate, latent_dim])
model = CVAE(latent_dim)
## 定义生成和保存图片的函数，
def generate_and_save_images(model, epoch, test_input):
    predictions = model.sample(test_input)
    fig = plt.figure(figsize=(4,4))
    for i in range(predictions.shape[0]):
        plt.subplot(4, 4, i+1)
        plt.imshow(predictions[i, :, :, 0], cmap='gray')
        plt.axis('off')
        # tight_layout minimizes the overlap between 2 sub-plots
    plt.savefig('chapter15/CVAE/image_at_epoch_{:04d}.png'.format(epoch))
    #plt.show()

## 生成初始图片
generate_and_save_images(model, 0, random_vector_for_generation)
## 定义优化算子
optimizer=tf.keras.optimizers.Adam(1e-4)
## 训练
for epoch in range(1, epochs + 1):
    start_time = time.time()
    for train_x in train_dataset:
        compute_apply_gradients(model, train_x, optimizer)
    end_time = time.time()
    if epoch % 1 == 0:
        loss = tf.keras.metrics.Mean()
        for test_x in test_dataset:
            loss(vae_kl_loss(model, test_x))
        elbo = -loss.result()
        display.clear_output(wait=False)
        print('Epoch: {}, Test set ELBO: {}, '
              'time elapse for current epoch {}'.format(epoch,
                                                        elbo,
                                                        end_time - start_time))
        ## 每个epoch，用decoder计算随机生成的向量，得到每个epoch生成的图片
        generate_and_save_images(
            model, epoch, random_vector_for_generation)

3.2 epochs=50两种损失函数生成的图片

两次试验损失函数包含相同的 Reconstruction Loss。