深度学习系列27：VAE生成模型

1. AE

AE（Autoencoder），自动编码器。自编码器的初衷是为了数据降维，假设原始特征x维度过高，那么我们希望通过编码器E将其编码成低维特征向量z=E(x)，编码的原则是尽可能保留原始信息，因此我们再训练一个解码器D，希望能通过z重构原始信息，即x≈D(E(x))，其优化目标一般是

我们常用的encoder-decoder即为最简单的一种AE。训练过程中加上一些扰动，就可以变成去噪自编码器（DAE）：

或者用遮盖（MIM，mask image modeling）的方法来加扰动：

2. VAE

损失为重构误差+KL散度。
对于每一个样本，需要用神经网络拟合均值$u$和方差${\delta }^2$，然后用标准正态分布采样得到Z，然后再恢复成X。其中方差项是核心，是用来进行对抗生成的关键。重构部分误差项会让$u$尽量接近真实值，而KL散度会保证生成器拥有标准的随机性。
我们来看下keras版本代码：

from keras.layers import Input, Dense, Lambda
from keras.models import Model
from keras import backend as K
# 输入
x = Input(shape=(original_dim,))
h = Dense(intermediate_dim, activation='relu')(x)

# 全连接层计算p(Z|X)的均值和方差
z_mean = Dense(latent_dim)(h)
z_log_var = Dense(latent_dim)(h)

def sampling(args):
    z_mean, z_log_var = args
    epsilon = K.random_normal(shape=K.shape(z_mean))
    return z_mean + K.exp(z_log_var / 2) * epsilon

# 使用均值和方差生成编码结果
z = Lambda(sampling, output_shape=(latent_dim,))([z_mean, z_log_var])

# 使用解码器恢复数据
decoder_h = Dense(intermediate_dim, activation='relu')
decoder_mean = Dense(original_dim, activation='sigmoid')
h_decoded = decoder_h(z)
x_decoded_mean = decoder_mean(h_decoded)

# 建立模型
vae = Model(x, x_decoded_mean)

# xent_loss是重构loss，kl_loss是KL loss
xent_loss = K.sum(K.binary_crossentropy(x, x_decoded_mean), axis=-1)
kl_loss = - 0.5 * K.sum(1 + z_log_var - K.square(z_mean) - K.exp(z_log_var), axis=-1)
vae_loss = K.mean(xent_loss + kl_loss)
vae.add_loss(vae_loss)
vae.compile(optimizer='rmsprop')
vae.summary()

# 开始训练
vae.fit(x_train,
        shuffle=True,
        epochs=epochs,
        batch_size=batch_size,
        validation_data=(x_test, None))

3. CVAE

CVAE是有标签下的条件VAE，最简单的方法，就是通过控制均值来控制生成图像的类别，只需要修改KL散度，从

改成

4. 直观解释

参见：https://spaces.ac.cn/archives/7725
首先是AE的解释，样本点经过降维后还要保证能够充分表达原信息，因此必须是在降维空间中分布非常规整无冗余的，下图中显然第四种是最好的。

变分自编码器通过加入后验分布p(z|x)，可以让编码空间更加“规整”。

现在每个样本x都对应一个“椭圆”，而确定一个“椭圆”需要两个信息：椭圆中心、椭圆轴长，它们各自构成一个向量，并且这个向量依赖于样本x，我们将其记为μ(x),σ(x)。既然整个椭圆都对应着样本x，我们要求椭圆内任意一点都可以重构x，所以训练目标为：

这样一来，VAE能达到两个效果：1、从标准高斯分布（单位圆）随机采样一个向量，就可以由解码器得到真实样本，即实现了生成模型；2、由于编码空间的紧凑形以及训练时对编码向量所加入的噪声，使得编码向量的各个分量能做到一定程度的解耦，并赋予编码向量一定的线性运算性质。

5. VQ-VAE

可以参考这篇：https://zhuanlan.zhihu.com/p/91434658
VQ-VAE的基础是逐点回归生成图像模型。由于图像像素点非常多，因此需要先降维，然后再用PixelCNN建模。降维过程如下：

具体来说，$n\ast n\ast 3$的图片被encode成$m\ast m\ast d$的矩阵，然后每个d维向量进行最近邻搜素，得到一个m*m的整数矩阵。一般来说，现在的m×m比原来的n×n×3要小得多，比如用CelebA数据做实验的时候，原来128×128×3的图可以编码为32×32的编码而基本不失真。
优化目标函数为：

前向计算的时候用的是$z_q$，计算梯度时用的是$z$

下面是使用keras的示例代码：

# 编码器
x_in = Input(shape=(img_dim, img_dim, 3))
......# 这里加入了一堆Conv2D, resnet层
e_model = Model(x_in, x)

# 解码器
z_in = Input(shape=K.int_shape(x)[1:])
......# 这里加入了一堆Conv2D, resnet层
g_model = Model(z_in, z)

# 硬编码模型（寻找最近邻）
z_in = Input(shape=K.int_shape(x)[1:])
z = z_in
vq_layer = VectorQuantizer(num_codes)
codes, code_vecs = vq_layer(z)
q_model = Model(z_in, [codes, code_vecs])

# 训练流程
z = e_model(x)
_, e = q_model(z) # 注意这里仅使用code_vec，没有用code。
ze = Lambda(lambda x: x[0] + K.stop_gradient(x[1] - x[0]))([z, e])
x = g_model(ze)

train_model = Model(x_in, [x, _])
mse_x = K.mean((x_in - x)**2)
mse_e = K.mean((K.stop_gradient(z) - e)**2) # 这里的e就是上文中的q_z.
mse_z = K.mean((K.stop_gradient(e) - z)**2)
loss = mse_x + mse_e + 0.25 * mse_z
train_model.add_loss(loss)
train_model.compile(optimizer=Adam(1e-3))

这里生成模型e_model使用了8个卷积层，其中前面2个缩小尺寸，后面6个加残差。尺寸由(None,128,128,3)变为(None, 32, 32, 128)
离散化的硬编码模型q_model是一个VectorQuantizer，需要特别看一下。计算inputs和embeddings之间的距离矩阵distance，尺寸为[2, 32, 32, 64]，然后用argmin得到下标，然后再用gather函数得到实际的向量。
解码模型g_model使用了3个卷积层（都有残差）+2个反卷积层（用于扩大尺寸）。尺寸由(None, 32, 32, 128)变为(None, 128, 128, 3)
总的模型是e_model->q_model(要加stop_gradient，即反向时消失)->g_model

class VectorQuantizer(Layer):
    """量化层
    """
    def __init__(self, num_codes, **kwargs):
        super(VectorQuantizer, self).__init__(**kwargs)
        self.num_codes = num_codes
        
    def build(self, input_shape): # 定义权重
        super(VectorQuantizer, self).build(input_shape)
        dim = input_shape[-1]
        self.embeddings = self.add_weight(
            name='embeddings',
            shape=(self.num_codes, dim),
            initializer='uniform'
        ) #使用add_weight函数创建权重矩阵
        
    def call(self, inputs): # 定义计算图
        """inputs.shape=[None, m, m, dim]
        """
        l2_inputs = K.sum(inputs**2, -1, keepdims=True)
        l2_embeddings = K.sum(self.embeddings**2, -1)
        for _ in range(K.ndim(inputs) - 1):
            l2_embeddings = K.expand_dims(l2_embeddings, 0)
        embeddings = K.transpose(self.embeddings)
        dot = K.dot(inputs, embeddings)
        distance = l2_inputs + l2_embeddings - 2 * dot
        codes = K.cast(K.argmin(distance, -1), 'int32')
        code_vecs = K.gather(self.embeddings, codes)
        return [codes, code_vecs]
    def compute_output_shape(self, input_shape):
        return [input_shape[:-1], input_shape]