 在 TensorFlow.org 上查看 |
 在 Google Colab 中运行 |
 在 GitHub 上查看源代码 |
 下载笔记本 |
此笔记本演示了如何基于 MNIST 数据集训练变分自编码器 (VAE) (1, 2)。VAE 是一种自编码器取值的概率分布，该模型会获取高维输入数据并将其压缩为较小的表示。与将输入映射到隐向量的传统自编码器不同，VAE 会将输入数据映射到概率分布的参数中，例如高斯分布的均值和方差。这种方式可以生成一个连续、结构化的隐空间，对于图像生成而言十分适用。
导入 Tensorflow 与其他库
pip install tensorflow-probability# to generate gifspip install imageiopip install git+https://github.com/tensorflow/docs
from IPython import display
import glob
import imageio
import matplotlib.pyplot as plt
import numpy as np
import PIL
import tensorflow as tf
import tensorflow_probability as tfp
import time
2023-11-07 19:07:43.934263: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered 2023-11-07 19:07:43.934314: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered 2023-11-07 19:07:43.935832: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
加载 MNIST 数据集
每个 MNIST 图像最初都是一个由 784 个整数组成的向量，每个整数在 0-255 之间，代表一个像素的强度。在我们的模型中使用伯努利分布对每个像素进行建模，并对数据集进行静态二值化。
(train_images, _), (test_images, _) = tf.keras.datasets.mnist.load_data()
Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz 11490434/11490434 [==============================] - 0s 0us/step
def preprocess_images(images):
images = images.reshape((images.shape[0], 28, 28, 1)) / 255.
return np.where(images > .5, 1.0, 0.0).astype('float32')
train_images = preprocess_images(train_images)
test_images = preprocess_images(test_images)
train_size = 60000
batch_size = 32
test_size = 10000
使用 tf.data 来将数据分批和打乱
train_dataset = (tf.data.Dataset.from_tensor_slices(train_images)
.shuffle(train_size).batch(batch_size))
test_dataset = (tf.data.Dataset.from_tensor_slices(test_images)
.shuffle(test_size).batch(batch_size))
通过 tf.keras.Sequential 连接生成网络与推理网络
在此 VAE 示例中，对编码器和解码器网络使用两个小型 ConvNet。在文献中，这些网络也分别称为推断/识别和生成模型。使用 tf.keras.Sequential 来简化实现。在下面的描述中，使 \(x\) 和 \(z\) 分别表示观测值和隐变量。
生成网络
这定义了近似后验分布 \(q(z|x)\)，它会将输入取作观测值并输出一组参数，用于指定隐变量表示 \(z\) 的条件分布。在本例中，简单地将分布建模为对角高斯分布，网络会输出分解高斯分布的均值和对数方差参数。输出对数方差而不是直接用于数值稳定性的方差。
推理网络
这定义了观测值的条件分布 \(p(x|z)\)，它会将隐变量样本 \(z\) 取作输入并输出观测值条件分布的参数。将隐变量先验分布 \(p(z)\) 建模为单位高斯分布。
重参数化技巧
要在训练期间为解码器生成样本 \(z\)，您可以在给定输入观测值 \(x\) 的情况下从编码器输出的参数所定义的隐变量分布中采样。然而，这种采样操作会产生瓶颈，因为反向传播不能流经随机节点。
要解决这个问题，请使用重参数化技巧。在我们的示例中，使用解码器参数和另一个参数 \(\epsilon\) 来逼近 \(z\)，如下所示：
\[z = \mu + \sigma \odot \epsilon\]
其中 \(\mu\) 和 \(\sigma\) 分别代表高斯分布的均值和标准差。它们可通过解码器输出推导得出。\(\epsilon\) 可被认为是用于保持 \(z\) 的随机性的随机噪声。从标准正态分布生成 \(\epsilon\)。
隐变量 \(z\) 现在由 \(\mu\)、\(\sigma\) 和 \(\epsilon\) 的函数生成，这将使模型能够分别通过 \(\mu\) 和 \(\sigma\) 在编码器中反向传播梯度，同时通过 \(\epsilon\) 保持随机性。
网络架构
对于编码器网络，使用两个卷积层后接一个全连接层。在解码器网络中，通过使用一个全连接层后接三个卷积转置层（在某些背景下也称为反卷积层）来镜像此架构。请注意，通常的做法是在训练 VAE 时避免使用批量归一化，因为使用 mini-batch 导致的额外随机性可能会在提高采样随机性的同时加剧不稳定性。
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
定义损失函数和优化器
VAE 通过最大化边际对数似然的证据下界（ELBO）进行训练：
\[\log p(x) \ge \text{ELBO} = \mathbb{E}_{q(z|x)}\left[\log \frac{p(x, z)}{q(z|x)}\right].\]
在实践中，优化此期望的单样本蒙特卡罗估值：
\(\log p(x| z) + \log p(z) - \log q(z|x)\)，其中 \(z\) 从 \(q(z|x)\) 中采样。
注：您也可以分析计算 KL 项，但为了简单起见，您在此处会将三个项全部应用到蒙特卡罗 Estimator 中。
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 =