tensorflow2 卷积变分自编码器的实现与简单应用

旧城等待， 2022-12-10 10:22 341阅读 0赞

### 1 编写目的 ###

VAE ( Variational Autoencoders，变分自编码) 模型组合了神经网络和贝叶斯推理这两种最好的方法，是最酷的神经网络，已经成为无监督学习的流行方法之一。（摘自《TensorFlow 深度学习实战》 安东尼奥-古利；阿米塔-卡普尔。机械工业出版社）。

在上一篇博客 [VAE 模型基本原理简单介绍][VAE] 已经介绍基本原理，这里记录一下 VAE 模型的实现过程（必须声明：）

本次实验采用的是 notebook，可以是自己电脑上安装的 jupyter notebook，也可以使用自己云服务器安装的，也可以考虑使用谷歌提供的 [Colaboratory][]。

### 2 卷积 VAE 模型的实现与测试 ###

> 注：本节内容绝大部分是参考 [https://tensorflow.google.cn/tutorials/generative/cvae][https_tensorflow.google.cn_tutorials_generative_cvae] (推荐阅读英文版)，只是对一些部分做了一些补充解释

推荐直接下载这个文件，[下载地址][Link 1]。下载后请去除后面的 `.txt` ，因为这是一个 ipynb 文件。

如果方便的话，可以直接访问 [VAE demo][]，使用 谷歌提供的免费 GPU 。

不能访问的话直接下载到本地，再运行。

#### 2.1 安装依赖 ####

1.  确定使用的是 `tensorflow 2.x`
    
        !pip show tensorflow
    
    如果当前安装的不是 tensorflow 2.x 的话，请输入以下命令安装：
    
        !pip install tensorflow==2.3.0 -i https://pypi.tuna.tsinghua.edu.cn/simple
2.  安装 `imageio`
    
        !pip install imageio

#### 2.2 导入相关库与加载数据集 ####

import tensorflow as tf
    
    import os
    import time
    import numpy as np
    import glob
    import matplotlib.pyplot as plt
    import PIL
    import imageio
    
    from IPython import display
    
    (train_images, _), (test_images, _) = tf.keras.datasets.mnist.load_data()
    
    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')
    
    # 标准化图片到区间 [0., 1.] 内
    train_images /= 255.
    test_images /= 255.
    
    # 二值化
    train_images[train_images >= .5] = 1.
    train_images[train_images < .5] = 0.
    test_images[test_images >= .5] = 1.
    test_images[test_images < .5] = 0.
    
    # 使用 tf.data 来将数据分批和打乱
    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)

#### 2.3 VAE 模型 ####

准备工作做完了后，这里正式开始编写实现VAE模型。

关于 VAE 模型的基本原理请参考一下 [上一篇博客][VAE]。

这里简单回顾一下， VAE 模型可以理解为 **三部分**

*  编码网络（Encoder Network），也称 `推断网络` 。该 NN 用来生成隐变量的参数（隐变量由多个高斯分布组成）。对于隐变量  z z z，首先初始化时可以是标准高斯分布，然后通过这个 NN，通过不断计算后验概率  q ( z ∣ x ) q(z|x) q(z∣x) 来逐步确定高斯分布的参数（均值和方差）。
 *  隐变量(Latent Variable)。作为 Encoder 过程的产物，隐变量至少能够包含一些输入数据的信息（降维的作用），同时也应该具有生成类似数据的潜力。
 *  解码网络(Decoder Network)，也称 `生成网络`。该 NN 用于根据隐变量生成数据，我们希望它既有能力还原 encoder 的数据，同时还能根据数据特征生成一些输入样本中不包含的数据。

**卷积**  
这里使用两个小卷积层分别用于 Encoder 和 Decoder 。在文献中，这些网络也分别被称为 推断网络或识别网络 （inference/recognition）和生成模型。

**重参数化技巧** (Reparameterization Trick)

*  训练过程中，为了生成样本  z z z 以便于 decoder 操作，我们可以从 encoder 生成的分布中进行采样。但是，由于反向传播无法通过随机节点，因此此采样操作会产生瓶颈。
 *  为了解决这个问题，我们使用了一个重新参数化的技巧。我们使用 decoder 参数和另一个参数  ε \\varepsilon ε 近似z，如下所示：  
     z = μ + σ   ⊙   ε z = \\mu +\\sigma \\ \\odot \\ \\varepsilon z=μ\+σ ⊙ ε

其中  μ \\mu μ 和  σ \\sigma σ 表示高斯分布中的均值 和 标准差。它们可以从 decoder 输出中导出。可以认为  ε \\varepsilon ε 是用来保持  z z z 随机性的随机噪声。我们从标准正态分布生成。

*  现在的  z z z 是  q ( z ∣ x ) q(z|x) q(z∣x) 生成（通过参数  μ \\mu μ， σ \\sigma σ 和  ε \\varepsilon ε），这将使模型分别通过  μ \\mu μ， σ \\sigma σ 在 encoder 中 反向传播梯度，同时通过  ε \\varepsilon ε 保持  z z z 的随机性。

**网络结构** (Network architecture)

对于 VAE 模型构建，

*  在 Encoder NN中，使用两个卷积层和一个完全连接的层。、
 *  在 Decoder NN中，通过使用一个完全连接的层和三个卷积转置层来镜像这种结构。

注意，在训练VAE时，通常避免使用批次标准化，因为使用小批量的额外随机性可能会加剧抽样随机性之外的不稳定性。

class CVAE(tf.keras.Model):
      """Convolutional variational autoencoder."""
    
      def __init__(self, latent_dim):
        super(CVAE, self).__init__()
        self.latent_dim = latent_dim
        self.encoder = 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.decoder = tf.keras.Sequential(
            [
                tf.keras.layers.InputLayer(input_shape=(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, padding='same',
                    activation='relu'),
                tf.keras.layers.Conv2DTranspose(
                    filters=32, kernel_size=3, strides=2, padding='same',
                    activation='relu'),
                # No activation
                tf.keras.layers.Conv2DTranspose(
                    filters=1, kernel_size=3, strides=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.encoder(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 * .5) + mean
    
      def decode(self, z, apply_sigmoid=False):
        logits = self.decoder(z)
        if apply_sigmoid:
          probs = tf.sigmoid(logits)
          return probs
        return logits

#### 2.4 定义损失函数和优化器 ####

上篇博客中提到过，VAE 通过  log ⁡ p ( x ) \\log p(x) logp(x) 极大似然 ELBO ( the evidence lower bound) 进行训练：  
 log ⁡ p ( x ) ≥ ELBO = E q ( z ∣ x ) \[ log ⁡ p ( x , z ) q ( z ∣ x ) \] . \\log p(x) \\ge \\text\{ELBO\} = \\mathbb\{E\}\_\{q(z|x)\}\\left\[\\log \\frac\{p(x, z)\}\{q(z|x)\}\\right\]. logp(x)≥ELBO=Eq(z∣x)\[logq(z∣x)p(x,z)\].

实际操作中，我们优化了这种单样本蒙特卡罗估计：  
 log ⁡ p ( x ∣ z ) + log ⁡ p ( z ) − log ⁡ q ( z ∣ x ) , \\log p(x| z) + \\log p(z) - \\log q(z|x), logp(x∣z)\+logp(z)−logq(z∣x),  
其中  z z z 从  q ( z ∣ x ) q(z|x) q(z∣x) 中采样。

optimizer = tf.keras.optimizers.Adam(1e-4)
    
    
    def log_normal_pdf(sample, mean, logvar, raxis=1):
      log2pi = tf.math.log(2. * np.pi)
      return tf.reduce_sum(
          -.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi),
          axis=raxis)
    
    
    def compute_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.)
      logqz_x = log_normal_pdf(z, mean, logvar)
      return -tf.reduce_mean(logpx_z + logpz - logqz_x)
    
    
    @tf.function
    def train_step(model, x, optimizer):
      """Executes one training step and returns the loss. This function computes the loss and gradients, and uses the latter to update the model's parameters. """
      with tf.GradientTape() as tape:
        loss = compute_loss(model, x)
      gradients = tape.gradient(loss, model.trainable_variables)
      optimizer.apply_gradients(zip(gradients, model.trainable_variables))

#### 2.5 训练模型与生成图片 ####

**训练**

*  我们从迭代数据集开始。
 *  在每次迭代期间，我们将图像传递给编码器，以获得近似后验  q ( z ∣ x ) q(z|x) q(z∣x) 的一组均值和对数方差参数(log-variance parameters）。
 *  然后，我们应用 *重参数化技巧* 从  q ( z ∣ x ) q(z|x) q(z∣x) 中采样。
 *  最后，我们将重新参数化的样本传递给解码器，以获取生成分布  p ( x ∣ z ) p(x|z) p(x∣z) 的 logit。
 *  **注意**：由于我们使用的是由 keras 加载的数据集，其中训练集中有 6 万个数据点，测试集中有 1 万个数据点，因此我们在测试集上的最终 ELBO 略高于对 Larochelle 版 MNIST 使用动态二值化的文献中的报告结果。
    
    > 关于 logits 的 [解释][Link 2]

**生成图片**

*  进行训练后，可以生成一些图片了。
 *  我们首先从单位高斯先验分布  p ( z ) p(z) p(z) 中采样一组潜在向量。
 *  随后生成器将潜在样本  z z z 转换为观测值的 logit，得到分布  p ( x ∣ z ) p(x|z) p(x∣z)。
 *  这里我们画出伯努利分布的概率。

epochs = 100
    # set the dimensionality of the latent space to a plane for visualization later
    latent_dim = 2
    num_examples_to_generate = 16
    
    # keeping the random vector constant for generation (prediction) so
    # it will be easier to see the improvement.
    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_sample):
      mean, logvar = model.encode(test_sample)
      z = model.reparameterize(mean, logvar)
      predictions = model.sample(z)
      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('image_at_epoch_{:04d}.png'.format(epoch))
      plt.show()
    
    # Pick a sample of the test set for generating output images
    assert batch_size >= num_examples_to_generate
    for test_batch in test_dataset.take(1):
      test_sample = test_batch[0:num_examples_to_generate, :, :, :]
    generate_and_save_images(model, 0, test_sample)
    
    for epoch in range(1, epochs + 1):
      start_time = time.time()
      for train_x in train_dataset:
        train_step(model, train_x, optimizer)
      end_time = time.time()
    
      loss = tf.keras.metrics.Mean()
      for test_x in test_dataset:
        loss(compute_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))
      generate_and_save_images(model, epoch, test_sample)

一百次循环后，生成的图片如下(感觉还可以)：

![在这里插入图片描述][2020101009492281.png_pic_center]  
同时可以生化 gif 图片来方便查看生成过程。

anim_file = 'cvae.gif'
    
    with imageio.get_writer(anim_file, mode='I') as writer:
      filenames = glob.glob('image*.png')
      filenames = sorted(filenames)
      for filename in filenames:
        image = imageio.imread(filename)
        writer.append_data(image)
      image = imageio.imread(filename)
      writer.append_data(image)

展示 gif 图片

import tensorflow_docs.vis.embed as embed
    embed.embed_file(anim_file)

![在这里插入图片描述][20201010101327220.gif_pic_center]

#### 2.6 继续玩隐变量 ####

隐变量

![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3NtaWxleWFuOQ_size_16_color_FFFFFF_t_70_pic_center]

#### 2.7 下一步工作 ####

以上内容绝大部分是摘录于tensorflow 官网的 [tutorials][] ，这里实现的是卷积 VAE 模型，简称(cVAE)。根据官网该文章， Next Step

*  应该通过增大网络大小来优化输出效果。比如说可以把 `Conv2D` 和 `Conv2DTranspose` 的层数调成 512等等 ；
 *  可以使用不同的数据集查看效果，比如 CIFAR-10

### 3 总结 ###

这里只是介绍一下如何快速体验一下 VAE 模型，模型的实现也是 tensorflow 官网提供的，对原来的文章做了简单的翻译工作，并且也加了一些自己的理解和补充。如果有更多好玩的实现demo，也将补充在这里。

> Smileyan  
> 2020.10.10 15:31

[VAE]: https://blog.csdn.net/smileyan9/article/details/107362252
[Colaboratory]: https://colab.research.google.com/notebooks/welcome.ipynb
[https_tensorflow.google.cn_tutorials_generative_cvae]: https://tensorflow.google.cn/tutorials/generative/cvae
[Link 1]: https://www.lanzoui.com/iFTFOhb35vc
[VAE demo]: https://colab.research.google.com/drive/1z-FjyjylVqtqol8xnhjqtP19gcVs1lLN?usp=sharing
[Link 2]: https://www.zhihu.com/question/60751553
[2020101009492281.png_pic_center]: /images/20221123/66ece861b06e4dce965d32c33efa0ea6.png
[20201010101327220.gif_pic_center]: /images/20221123/b5da4c34aedd49b9a57dd40cd58123b4.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3NtaWxleWFuOQ_size_16_color_FFFFFF_t_70_pic_center]: /images/20221123/fa139d2d1e394b48ab0eeda4a5eed44c.png
[tutorials]: https://www.tensorflow.org/tutorials/generative/cvae