TF2 VAE原理及代碼實現
VAE（Variational Autoencoder）
相關文章導航
損失函數
KL散度
相對熵（relative entropy），又稱爲KL散度（Kullback - Leibler divergence）
，信息散度，信息增益（information gain）。
KL散度是兩個概率分佈P和Q差別的非對稱性的度量
KL散度是用來度量使用基於Q的編碼來編碼來自P的樣本平均所需的額外的比特個數。 典型情況下，P表示數據的真實分佈，Q表示數據的理論分佈，模型分佈，或P的近似分佈。
根據shannon的信息論，給定一個字符集的概率分佈，我們可以設計一種編碼，使得表示該字符集組成的字符串平均需要的比特數最少。假設這個字符集是X，對x∈X，其出現概率爲P(x)，那麼其最優編碼平均需要的比特數等於這個字符集的熵：
在同樣的字符集上，假設存在另一個概率分佈Q(X)。如果用概率分佈P(X)的最優編碼（即字符x的編碼長度等於log[1/P(x)]），來爲符合分佈Q(X)的字符編碼，那麼表示這些字符就會比理想情況多用一些比特數。KL-divergence就是用來衡量這種情況下平均每個字符多用的比特數，因此可以用來衡量兩個分佈的距離。即：
由於-log(u)是凸函數，因此有下面的不等式
即KL-divergence始終是大於等於0的。當且僅當兩分佈相同時，KL - divergence等於0。
KL散度是兩個概率分佈P和Q差別的非對稱性的度量。KL散度是用來度量使用基於Q的編碼來編碼來自P的樣本平均所需的額外的位元數。
所以KL散度被運用在VAE中，它可以用來量度潛在變量的分佈和單位高斯分佈的差異。 被設計在了損失函數中。
但是KL散度在VAE的運用進行講話，一下的公式的推導過程
tf代碼表示損失函數
rec_loss = tf.nn.sigmoid_cross_entropy_with_logits(labels=x, logits=x_rec_logits)
rec_loss = tf.reduce_sum(rec_loss) / x.shape[0]
# compute kl divergence (mu, var) ~ N (0, 1)
# https://stats.stackexchange.com/questions/7440/kl-divergence-between-two-univariate-gaussians
kl_div = -0.5 * (log_var + 1 - mu ** 2 - tf.exp(log_var))
kl_div = tf.reduce_sum(kl_div) / x.shape[0]
代碼
在知道了VAE的損失函數後我們就可以開始逐步的更新網絡接口來寫代碼了
數據集
數據集的我們選擇tensorflow中自帶的fashion_minist通過
tf.keras.datasets.fashion_mnist.load_data() 來加載即可
下面開始正式的寫代碼了
導入相關的模塊
每次導入的模塊都差不多，可以把這些代碼固定在剪貼板上之後的時候直接複製過來就好
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import tensorflow as tf
import numpy as np
from tensorflow import keras
from tensorflow.keras import Sequential, layers
from PIL import Image
from matplotlib import pyplot as plt
from tensorflow.compat.v1 import ConfigProto
from tensorflow.compat.v1 import InteractiveSession
config = ConfigProto()
config.gpu_options.allow_growth = True
session = InteractiveSession(config=config)
tf.random.set_seed(2233)
np.random.seed(2233)
assert tf.__version__.startswith('2.')
定義一個保存圖片的函數
爲了看到網絡訓練的效果我們寫一個拼接多張圖片並保存到本地的一個函數
def save_images(imgs, name):
new_im = Image.new('L', (280, 280))
index = 0
for i in range(0, 280, 28):
for j in range(0, 280, 28):
im = imgs[index]
im = Image.fromarray(im, mode='L')
new_im.paste(im, (i, j))
index += 1
new_im.save(name)
加載數據集
這一部分沒啥好說的，都快寫臭了，但是要注意的是我們在構建數據集的時候不需要lable，只要給圖片數據就可以了。
h_dim = 20
batchsz = 512
lr = 1e-3
(x_train, y_train), (x_test, y_test) = keras.datasets.fashion_mnist.load_data()
x_train, x_test = x_train.astype(np.float32) / 255., x_test.astype(np.float32) / 255
# we do not need label
train_db = tf.data.Dataset.from_tensor_slices(x_train)
train_db = train_db.shuffle(batchsz * 5).batch(batchsz)
test_db = tf.data.Dataset.from_tensor_slices(x_test)
test_db = test_db.batch(batchsz)
print(x_train.shape, y_train.shape)
print(x_test.shape, y_test.shape)
z_dim = 10
搭建網絡
根據這個圖片來搭建網絡，整體的搭建還是很簡單的屬於一個比較小型的網絡，只有幾十萬個參數。
class VAE(keras.Model):
def __init__(self):
super(VAE, self).__init__()
# Encoders
self.fc1 = layers.Dense(128)
self.fc2 = layers.Dense(z_dim)
self.fc3 = layers.Dense(z_dim)
# Decoders
self.fc4 = layers.Dense(128)
self.fc5 = layers.Dense(784)
def encoder(self, x):
h = tf.nn.relu(self.fc1(x))
mu = self.fc2(h)
# get variance
log_var = self.fc3(h)
return mu, log_var
def decoder(self, z):
out = tf.nn.relu(self.fc4(z))
out = self.fc5(out)
return out
def reparameterize(self, mu, log_var):
eps = tf.random.normal(log_var.shape)
std = tf.exp(log_var) ** 0.5
z = mu + std * eps
return z
def call(self, inputs, training=None):
# [b,784] => [b,z_dim]
mu, log_var = self.encoder(inputs)
# reparameterization trick
z = self.reparameterize(mu, log_var)
x_hat = self.decoder(z)
return x_hat, mu, log_var
模型訓練
模型在訓練的時候，需要注意的就是損失函數的書寫，這個在上面已經寫過了我們就直接上代碼了
model = VAE()
model.build(input_shape=(4, 784))
model.summary()
optimizer = tf.optimizers.Adam(lr=lr)
for epoch in range(1000):
for step, x in enumerate(train_db):
# [b,28,28] => [b,784]
x = tf.reshape(x, [-1, 784])
with tf.GradientTape() as tape:
x_hat, mu, log_var = model(x)
rec_loss = tf.losses.binary_crossentropy(x, x_hat, from_logits=True)
rec_loss = tf.reduce_mean(rec_loss)
# compute kl divergence (mu, var) ~ N(0,1)
kl_div = -0.5 * (log_var + 1 - mu ** 2 - tf.exp(log_var))
kl_div = tf.reduce_mean(kl_div) / x.shape[0]
loss = rec_loss + 1. * kl_div
grads = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(grads, model.trainable_variables))
if step % 100 == 0:
print(epoch, step, 'kl div', float(kl_div), 'rec loss', float(rec_loss))
# evaluation
z = tf.random.normal((batchsz, z_dim))
logits = model.decoder(z)
x_hat = tf.sigmoid(logits)
x_hat = tf.reshape(x_hat, [-1, 28, 28]).numpy() * 255.
x_hat = x_hat.astype(np.uint8)
save_images(x_hat, 'vae_imgs/sampled_epoch_%d.png' % epoch)
訓練效果
在測試效果的時候我們使用的是完全隨機的數字，下面讓我們來看看效果怎麼樣
第一張 訓練一個epoch
第二張 訓練兩個epoch
第三張 訓練三個epoch
第6張， 訓練6個epoch
第15張 訓練15個epoch
第20張 訓練20個epoch
第100 張， 訓練100個epoch
可以看到這時候的效果已經不錯了哦
第150張 第150個epoch
完整代碼
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import tensorflow as tf
import numpy as np
from tensorflow import keras
from tensorflow.keras import Sequential, layers
from PIL import Image
from matplotlib import pyplot as plt
from tensorflow.compat.v1 import ConfigProto
from tensorflow.compat.v1 import InteractiveSession
config = ConfigProto()
config.gpu_options.allow_growth = True
session = InteractiveSession(config=config)
tf.random.set_seed(2233)
np.random.seed(2233)
assert tf.__version__.startswith('2.')
def save_images(imgs, name):
new_im = Image.new('L', (280, 280))
index = 0
for i in range(0, 280, 28):
for j in range(0, 280, 28):
im = imgs[index]
im = Image.fromarray(im, mode='L')
new_im.paste(im, (i, j))
index += 1
new_im.save(name)
h_dim = 20
batchsz = 512
lr = 1e-3
(x_train, y_train), (x_test, y_test) = keras.datasets.fashion_mnist.load_data()
x_train, x_test = x_train.astype(np.float32) / 255., x_test.astype(np.float32) / 255
# we do not need label
train_db = tf.data.Dataset.from_tensor_slices(x_train)
train_db = train_db.shuffle(batchsz * 5).batch(batchsz)
test_db = tf.data.Dataset.from_tensor_slices(x_test)
test_db = test_db.batch(batchsz)
print(x_train.shape, y_train.shape)
print(x_test.shape, y_test.shape)
z_dim = 10
class VAE(keras.Model):
def __init__(self):
super(VAE, self).__init__()
# Encoders
self.fc1 = layers.Dense(128)
self.fc2 = layers.Dense(z_dim)
self.fc3 = layers.Dense(z_dim)
# Decoders
self.fc4 = layers.Dense(128)
self.fc5 = layers.Dense(784)
def encoder(self, x):
h = tf.nn.relu(self.fc1(x))
mu = self.fc2(h)
# get variance
log_var = self.fc3(h)
return mu, log_var
def decoder(self, z):
out = tf.nn.relu(self.fc4(z))
out = self.fc5(out)
return out
def reparameterize(self, mu, log_var):
eps = tf.random.normal(log_var.shape)
std = tf.exp(log_var) ** 0.5
z = mu + std * eps
return z
def call(self, inputs, training=None):
# [b,784] => [b,z_dim]
mu, log_var = self.encoder(inputs)
# reparameterization trick
z = self.reparameterize(mu, log_var)
x_hat = self.decoder(z)
return x_hat, mu, log_var
model = VAE()
model.build(input_shape=(4, 784))
model.summary()
optimizer = tf.optimizers.Adam(lr=lr)
for epoch in range(1000):
for step, x in enumerate(train_db):
# [b,28,28] => [b,784]
x = tf.reshape(x, [-1, 784])
with tf.GradientTape() as tape:
x_hat, mu, log_var = model(x)
rec_loss = tf.losses.binary_crossentropy(x, x_hat, from_logits=True)
rec_loss = tf.reduce_mean(rec_loss)
# compute kl divergence (mu, var) ~ N(0,1)
kl_div = -0.5 * (log_var + 1 - mu ** 2 - tf.exp(log_var))
kl_div = tf.reduce_mean(kl_div) / x.shape[0]
loss = rec_loss + 1. * kl_div
grads = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(grads, model.trainable_variables))
if step % 100 == 0:
print(epoch, step, 'kl div', float(kl_div), 'rec loss', float(rec_loss))
# evaluation
z = tf.random.normal((batchsz, z_dim))
logits = model.decoder(z)
x_hat = tf.sigmoid(logits)
x_hat = tf.reshape(x_hat, [-1, 28, 28]).numpy() * 255.
x_hat = x_hat.astype(np.uint8)
save_images(x_hat, 'vae_imgs/sampled_epoch_%d.png' % epoch)
參考書籍: TensorFlow 深度學習 — 龍龍老師