「TensorFlow」變分自編碼器實現

知識 01-19

摘要：一個小例子帶你玩轉(TensorFlow)變分自編碼器！

變分自編碼器（VAEs）是學習低維數據表示的強大模型。TensorFlow的分發包提供了一種簡單的方法來實現不同類型的VAE。

在這篇文章中，我將引導你完成在MNIST上訓練簡單VAE的步驟，主要側重於實戰。

1.定義網路：

VAE由三部分組成：編碼器q（z | x ），先驗p（z ），解碼器p（x | z ）。

「TensorFlow」變分自編碼器實現

編碼器將圖像映射到針對該圖像的代碼的分布上。這種分布也被稱為後驗（posterior），因為它反映了我們關於代碼應該用於給定圖像之後的準確度。

import tensorflow as tf
tfd = tf.contrib.distributionsdef make_encoder(data, code_size):
x = tf.layers.flatten(data)
x = tf.layers.dense(x, 200, tf.nn.relu)
x = tf.layers.dense(x, 200, tf.nn.relu)
loc = tf.layers.dense(x, code_size)
scale = tf.layers.dense(x, code_size, tf.nn.softplus)
return tfd.MultivariateNormalDiag(loc, scale)

先驗（prior）是固定的，上面我們還定義了我們期望的代碼塊分布。這對VAE可以使用的代碼塊提供了一個非常彈性的選擇空間。它通常只是一個零均值和單位方差的正態分布。

def make_prior(code_size):
loc = tf.zeros(code_size)
scale = tf.ones(code_size)
return tfd.MultivariateNormalDiag(loc, scale)

解碼器需要編寫一個代碼塊並將其映射回合理的圖像分布。它允許我們重建圖像，或為我們選擇的任何代碼塊生成新圖像。

import numpy as np
def make_decoder(code, data_shape):
x = code
x = tf.layers.dense(x, 200, tf.nn.relu)
x = tf.layers.dense(x, 200, tf.nn.relu)
logit = tf.layers.dense(x, np.prod(data_shape))
logit = tf.reshape(logit, [-1] + data_shape)
return tfd.Independent(tfd.Bernoulli(logit), 2)

在這裡，我們對數據使用伯努利分布，將像素建模為二進位值。根據數據的類型和領域，您可能需要以不同的方式對其進行建模，例如再次以正態分布的形式進行建模。

該tfd.Independent(..., 2)告訴TensorFlow，內部兩個尺寸——（寬度和高度）。在我們的例子中，屬於同一個數據點，即使他們有獨立的參數。這使我們能夠評估分布下圖像的概率，而不僅僅是單個像素。

2.重新使用模型

我們希望使用解碼器網路兩次，計算下一節中描述的重構損失，以及一些隨機採樣的可視化代碼塊。

在TensorFlow中，如果您調用兩次網路功能，它將創建兩個獨立的網路。TensorFlow模板允許您打包一個函數，以便多次調用它將重用相同的網路參數。

make_encoder = tf.make_template("encoder", make_encoder)
make_decoder = tf.make_template("decoder", make_decoder)

之前沒有可訓練的參數，所以我們不需要將其包裝到模板中。

3.定義損失

我們希望找到為我們的數據集分配最高可能性的網路參數。然而，數據點的可能性取決於最好的代碼塊，這點在訓練中是我們不知道。

代替的，我們將使用（ELBO）對數據可能性進行近似訓練。

「TensorFlow」變分自編碼器實現

這裡的重要細節是，ELBO僅使用給定我們當前對其代碼塊的估計的數據點的可能性，我們可以對其進行抽樣。

data = tf.placeholder(tf.float32, [None, 28, 28])prior = make_prior(code_size=2)posterior = make_encoder(data, code_size=2)code = posterior.sample()likelihood = make_decoder(code, [28, 28]).log_prob(data)divergence = tfd.kl_divergence(posterior, prior)elbo = tf.reduce_mean(likelihood - divergence)

一個直觀的解釋是，最大化ELBO可以最大限度地提供給定當前代碼的數據的可能性，同時鼓勵代碼接近我們先前的代碼應該是什麼樣子的信念。

4.運行訓練

我們使用梯度下降來最大化ELBO。這是一個非常可行的方案，因為採樣操作是在內部使用重新參數化技巧來實現的，所以TensorFlow可以通過它們反向傳播。

optimize = tf.train.AdamOptimizer(0.001).minimize(-elbo)

而且，我們從之前的樣本中抽取一些隨機代碼來可視化VAE學到的相應圖像。這就是上面我們使用tf.make_template()的原因，讓我們再次調用解碼器網路。

samples = make_decoder(prior.sample(10), [28, 28]).mean()

最後，我們載入數據並創建一個會話來運行訓練：

from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets("MNIST_data/")with tf.train.MonitoredSession() as sess:
for epoch in range(20):
test_elbo, test_codes, test_samples = sess.run(
[elbo, code, samples], {data: mnist.test.images})
print("Epoch", epoch, "elbo", test_elbo)
plot_codes(test_codes)
plot_sample(test_samples)
for _ in range(600):
sess.run(optimize, {data: mnist.train.next_batch(100)[0]})

如果你想玩代碼，看看完整的代碼示例，它也包含在這篇文章中其他省略的繪圖，以及前幾個訓練階段的解碼樣本。

「TensorFlow」變分自編碼器實現