variational_autoencoder.ipynb (Source)
Preamble¶
In [1]:
%matplotlib notebook
In [2]:
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import norm
from keras import backend as K
from keras.layers import (Input, InputLayer, Dense, Lambda, Layer,
Add, Multiply)
from keras.models import Model, Sequential
from keras.datasets import mnist
In [3]:
import pandas as pd
from matplotlib.ticker import FormatStrFormatter
from keras.utils.vis_utils import model_to_dot, plot_model
from IPython.display import SVG
Notebook Configuration¶
In [4]:
np.set_printoptions(precision=2,
edgeitems=3,
linewidth=80,
suppress=True)
In [5]:
'TensorFlow version: ' + K.tf.__version__
Out[5]:
Dataset (MNIST)¶
In [6]:
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train = np.expand_dims(x_train, axis=-1) / 255.
x_test = np.expand_dims(x_test, axis=-1) / 255.
In [7]:
img_rows, img_cols, img_chns = x_train.shape[1:]
Constant definitions¶
In [8]:
original_dim = img_rows * img_cols
intermediate_dim = 256
latent_dim = 2
batch_size = 100
epochs = 50
epsilon_std = 1.0
Model specification¶
Encoder¶
Inference network¶
In [9]:
x = Input(shape=(original_dim,), name='x')
h = Dense(intermediate_dim, activation='relu',
name='hidden_enc')(x)
z_mu = Dense(latent_dim, name='mu')(h)
z_log_var = Dense(latent_dim, name='log_var')(h)
z_sigma = Lambda(lambda t: K.exp(.5*t), name='sigma')(z_log_var)
In [10]:
m = Model(inputs=x, outputs=[z_mu, z_log_var])
In [11]:
SVG(model_to_dot(m, show_shapes=False)
.create(prog='dot', format='svg'))
Out[11]:
In [12]:
plot_model(
model=m, show_shapes=False,
to_file='../../images/vae/inference_network.svg'
)
Reparameterization with Merge Layers¶
In [13]:
z_mu = Input(shape=(latent_dim,), name='mu')
z_sigma = Input(shape=(latent_dim,), name='sigma')
eps = Input(shape=(latent_dim,), name='eps')
z_eps = Multiply(name='z_eps')([z_sigma, eps])
z = Add(name='z')([z_mu, z_eps])
In [14]:
m = Model(inputs=[eps, z_mu, z_sigma], outputs=z)
In [15]:
SVG(model_to_dot(m, show_shapes=False)
.create(prog='dot', format='svg'))
Out[15]:
In [16]:
plot_model(
model=m, show_shapes=False,
to_file='../../images/vae/reparameterization.svg'
)
Simplified architecture visualization¶
In [17]:
x = Input(shape=(original_dim,), name='x')
h = Dense(intermediate_dim, activation='relu',
name='hidden_enc')(x)
z_mu = Dense(latent_dim, name='mu')(h)
z_log_var = Dense(latent_dim, name='log_var')(h)
z_sigma = Lambda(lambda t: K.exp(.5*t), name='sigma')(z_log_var)
eps = Input(shape=(latent_dim,), name='eps')
z_eps = Multiply(name='z_eps')([z_sigma, eps])
z = Add(name='z')([z_mu, z_eps])
In [18]:
encoder = Model(inputs=[x, eps], outputs=z)
In [19]:
SVG(model_to_dot(encoder, show_shapes=False)
.create(prog='dot', format='svg'))
Out[19]:
In [20]:
plot_model(
model=encoder, show_shapes=False,
to_file='../../images/vae/encoder.svg'
)
In [21]:
plot_model(
model=encoder, show_shapes=True,
to_file='../../images/vae/encoder_shapes.svg'
)
Full architecture visualization with auxiliary layers¶
In [22]:
class KLDivergenceLayer(Layer):
""" Identity transform layer that adds KL divergence
to the final model loss.
"""
def __init__(self, *args, **kwargs):
self.is_placeholder = True
super(KLDivergenceLayer, self).__init__(*args, **kwargs)
def call(self, inputs):
mu, log_var = inputs
kl_batch = - .5 * K.sum(1 + log_var -
K.square(mu) -
K.exp(log_var), axis=-1)
self.add_loss(K.mean(kl_batch), inputs=inputs)
return inputs
In [23]:
x = Input(shape=(original_dim,), name='x')
h = Dense(intermediate_dim, activation='relu',
name='hidden_enc')(x)
z_mu = Dense(latent_dim, name='mu')(h)
z_log_var = Dense(latent_dim, name='log_var')(h)
z_mu, z_log_var = KLDivergenceLayer(name='kl')([z_mu, z_log_var])
z_sigma = Lambda(lambda t: K.exp(.5*t), name='sigma')(z_log_var)
eps = Input(tensor=K.random_normal(shape=(K.shape(x)[0],
latent_dim)), name='eps')
z_eps = Multiply(name='z_eps')([z_sigma, eps])
z = Add(name='z')([z_mu, z_eps])
In [24]:
encoder = Model(inputs=[x, eps], outputs=z)
In [25]:
SVG(model_to_dot(encoder, show_shapes=False)
.create(prog='dot', format='svg'))
Out[25]:
In [26]:
plot_model(
model=encoder, show_shapes=False,
to_file='../../images/vae/encoder_full.svg'
)
In [27]:
plot_model(
model=encoder, show_shapes=True,
to_file='../../images/vae/encoder_full_shapes.svg'
)
Decoder¶
In [28]:
decoder = Sequential([
Dense(intermediate_dim, input_dim=latent_dim,
activation='relu', name='hidden_dec'),
Dense(original_dim, activation='sigmoid', name='x_pred')
], name='decoder')
In [29]:
# equivalent to above. we specify InputLayer
# explicitly to set layer name for architecture diagram
decoder = Sequential([
InputLayer(input_shape=(latent_dim,), name='z'),
Dense(intermediate_dim, input_shape=(latent_dim,),
activation='relu', name='hidden_dec'),
Dense(original_dim, activation='sigmoid', name='x_pred')
], name='decoder')
In [30]:
SVG(model_to_dot(decoder, show_shapes=False)
.create(prog='dot', format='svg'))
Out[30]:
In [31]:
plot_model(
model=decoder, show_shapes=False,
to_file='../../images/vae/decoder.svg'
)
Specifying the VAE¶
In [32]:
x_pred = decoder(z)
In [33]:
# again, equivalent to above. fleshing it out fully
# for final end-to-end vae architecture visualization;
# otherwise, sequential models just get chunked into
# single layer
h_dec = Dense(intermediate_dim, activation='relu',
name='hidden_dec')(z)
x_pred = Dense(original_dim, activation='sigmoid',
name='x_pred')(h_dec)
In [34]:
vae = Model(inputs=[x, eps], outputs=x_pred)
In [35]:
SVG(model_to_dot(vae, show_shapes=True)
.create(prog='dot', format='svg'))
Out[35]:
In [36]:
plot_model(
model=vae, show_shapes=False,
to_file='../../images/vae/vae_full.svg'
)
In [37]:
plot_model(
model=vae, show_shapes=True,
to_file='../../images/vae/vae_full_shapes.svg'
)
Putting it all together¶
In [38]:
def nll(y_true, y_pred):
""" Negative log likelihood (Bernoulli). """
# keras.losses.binary_crossentropy gives the mean
# over the last axis. we require the sum
return K.sum(K.binary_crossentropy(y_true, y_pred), axis=-1)
In [39]:
x = Input(shape=(original_dim,))
h = Dense(intermediate_dim, activation='relu')(x)
z_mu = Dense(latent_dim)(h)
z_log_var = Dense(latent_dim)(h)
z_mu, z_log_var = KLDivergenceLayer()([z_mu, z_log_var])
z_sigma = Lambda(lambda t: K.exp(.5*t))(z_log_var)
eps = Input(tensor=K.random_normal(shape=(K.shape(x)[0],
latent_dim)))
z_eps = Multiply()([z_sigma, eps])
z = Add()([z_mu, z_eps])
decoder = Sequential([
Dense(intermediate_dim, input_dim=latent_dim, activation='relu'),
Dense(original_dim, activation='sigmoid')
])
x_pred = decoder(z)
In [40]:
vae = Model(inputs=[x, eps], outputs=x_pred, name='vae')
vae.compile(optimizer='rmsprop', loss=nll)
Model fitting¶
In [41]:
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train = x_train.reshape(-1, original_dim) / 255.
x_test = x_test.reshape(-1, original_dim) / 255.
In [42]:
hist = vae.fit(
x_train,
x_train,
shuffle=True,
epochs=epochs,
batch_size=batch_size,
validation_data=(x_test, x_test)
)
Model Evaluation¶
In [43]:
golden_size = lambda width: (width, 2. * width / (1 + np.sqrt(5)))
NELBO¶
In [44]:
fig, ax = plt.subplots(figsize=golden_size(6))
hist_df = pd.DataFrame(hist.history)
hist_df.plot(ax=ax)
ax.set_ylabel('NELBO')
ax.set_xlabel('# epochs')
ax.set_ylim(.99*hist_df[1:].values.min(),
1.1*hist_df[1:].values.max())
plt.savefig('../../images/vae/nelbo.svg', format='svg')
plt.show()
Observed space manifold¶
In [45]:
# display a 2D manifold of the images
n = 15 # figure with 15x15 images
quantile_min = 0.01
quantile_max = 0.99
# linearly spaced coordinates on the unit square were transformed
# through the inverse CDF (ppf) of the Gaussian to produce values
# of the latent variables z, since the prior of the latent space
# is Gaussian
z1 = norm.ppf(np.linspace(quantile_min, quantile_max, n))
z2 = norm.ppf(np.linspace(quantile_max, quantile_min, n))
z_grid = np.dstack(np.meshgrid(z1, z2))
In [46]:
x_pred_grid = decoder.predict(z_grid.reshape(n*n, latent_dim)) \
.reshape(n, n, img_rows, img_cols)
In [47]:
fig, ax = plt.subplots(figsize=(5, 5))
ax.imshow(np.block(list(map(list, x_pred_grid))), cmap='gray')
ax.set_xticks(np.arange(0, n*img_rows, img_rows) + .5 * img_rows)
ax.set_xticklabels(map('{:.2f}'.format, z1), rotation=90)
ax.set_yticks(np.arange(0, n*img_cols, img_cols) + .5 * img_cols)
ax.set_yticklabels(map('{:.2f}'.format, z2))
ax.set_xlabel('$z_1$')
ax.set_ylabel('$z_2$')
plt.savefig('../../images/vae/result_manifold.png')
plt.show()
In [48]:
# deterministic test time encoder
test_encoder = Model(x, z_mu)
# display a 2D plot of the digit classes in the latent space
z_test = test_encoder.predict(x_test, batch_size=batch_size)
In [49]:
fig, ax = plt.subplots(figsize=(6, 5))
cbar = ax.scatter(z_test[:, 0], z_test[:, 1], c=y_test,
alpha=.4, s=3**2, cmap='viridis')
fig.colorbar(cbar, ax=ax)
ax.set_xlim(2.*norm.ppf((quantile_min, quantile_max)))
ax.set_ylim(2.*norm.ppf((quantile_min, quantile_max)))
ax.set_xlabel('$z_1$')
ax.set_ylabel('$z_2$')
plt.savefig('../../images/vae/result_latent_space.png')
plt.show()
In [50]:
fig, (ax1, ax2) = plt.subplots(ncols=2, figsize=(12, 4.5))
ax1.imshow(np.block(list(map(list, x_pred_grid))), cmap='gray')
ax1.set_xticks(np.arange(0, n*img_rows, img_rows) + .5 * img_rows)
ax1.set_xticklabels(map('{:.2f}'.format, z1), rotation=90)
ax1.set_yticks(np.arange(0, n*img_cols, img_cols) + .5 * img_cols)
ax1.set_yticklabels(map('{:.2f}'.format, z2))
ax.set_xlabel('$z_1$')
ax.set_ylabel('$z_2$')
cbar = ax2.scatter(z_test[:, 0], z_test[:, 1], c=y_test,
alpha=.4, s=3**2, cmap='viridis')
fig.colorbar(cbar, ax=ax2)
ax2.set_xlim(2.*norm.ppf((quantile_min, quantile_max)))
ax2.set_ylim(2.*norm.ppf((quantile_min, quantile_max)))
ax2.set_xlabel('$z_1$')
ax2.set_ylabel('$z_2$')
plt.savefig('../../images/vae/result_combined.png')
plt.show()
