Autoencoding beyond pixels using a learned similarity metric
VAE with GAN Discriminator behind Use the feature map extracted from the middle layer of Discriminator instead of the image generated by Decoder for VAE error. While VAE measures the error in pixel units, blurry images appear, but by measuring the error in feature map units, it may be possible to generate fine images while reproducing global features.
Encoder learns the difference between the feature map of the original image and the decoded image as an error (VAE) In addition to that, Decoder learns the discrimination result of the decoded image / randomly generated image by Discriminator as an error (VAE + GAN). Discriminator learns from the identification results of original images, decoded images, and randomly generated images (GAN)
It seems good to pre-learn VAE with normal pixel units and GAN with original image and randomly generated image (replace GAN Generator with VAE Decoder and learn both independently)
I'm still experimenting, so I may change it later.
In the case of DCGAN, I tried my best to bring it to the original implementation, but this time I finished with a simplified version where each of Encoder, Decoder, Discriminator is trained with different samples
I referred to the original paper and the implementation of others, but made some changes.
--VAE and GAN were trained at the same time in Pretraining, but VAE → GAN are trained in order. --In Mean Square and Kullback Leibler Divergence, the dimensions of Feature and so on were summed (Sum), and the dimensions of Sample were averaged (Mean), but when I changed the image size and latent variables, the ratio changed and I felt it was troublesome. So all changed to Mean (may be mathematically wrong) --Decoder Loss was changed from Encoder Loss + Discriminator Loss to the average of both
vaegan.py
from __future__ import (
division,
print_function,
absolute_import
)
from six.moves import range
import tensorflow as tf
import tflearn
import os
import numpy as np
from skimage import io
PRE_VAE_TENSORBOARD_DIR = '/tmp/tflearn_logs/vae/'
PRE_DIS_TENSORBOARD_DIR = '/tmp/tflearn_logs/dis/'
PRE_VAE_CHECKPOINT_PATH = '/tmp/vaegan/pre-vae'
PRE_DIS_CHECKPOINT_PATH = '/tmp/vaegan/pre-dis'
CHECKPOINT_PATH = '/tmp/vaegan/model'
DNN = tflearn.DNN
input_data = tflearn.input_data
fc = tflearn.fully_connected
reshape = tflearn.reshape
conv = tflearn.conv_2d
conv_t = tflearn.conv_2d_transpose
max_pool = tflearn.max_pool_2d
bn = tflearn.batch_normalization
merge = tflearn.merge
sigmoid = tflearn.sigmoid
softmax = tflearn.softmax
softplus = tflearn.softplus
relu = tflearn.relu
elu = tflearn.elu
crossentropy = tflearn.categorical_crossentropy
adam = tflearn.Adam
Trainer = tflearn.Trainer
TrainOp = tflearn.TrainOp
if not os.path.exists('/tmp/tflearn_logs'):
os.mkdir('/tmp/tflearn_logs')
if not os.path.exists(PRE_VAE_TENSORBOARD_DIR):
os.mkdir(PRE_VAE_TENSORBOARD_DIR)
if not os.path.exists(PRE_DIS_TENSORBOARD_DIR):
os.mkdir(PRE_DIS_TENSORBOARD_DIR)
if not os.path.exists('/tmp/vaegan/'):
os.mkdir('/tmp/vaegan/')
class VAEGAN(object):
def __init__(self, img_shape, n_first_channel, n_layer, latent_dim,
kullback_leibler_ratio, reconstruction_weight_against_detail,
vae_learning_rate=0.001, vae_beta1=0.5,
discriminator_learning_rate=0.00001, discriminator_beta1=0.5):
self.img_shape = list(img_shape)
self.input_shape = [None] + self.img_shape
self.img_size = img_shape[:2]
self.n_first_channel = n_first_channel
self.n_layer = n_layer
self.kullback_leibler_ratio = kullback_leibler_ratio
self.reconstruction_weight_against_detail = reconstruction_weight_against_detail
self.latent_dim = latent_dim
self.vae_learning_rate = vae_learning_rate
self.vae_beta1 = vae_beta1
self.discriminator_learning_rate = discriminator_learning_rate
self.discriminator_beta1 = discriminator_beta1
assert self.n_layer > 1, 'n_layer must be more than 1'
self.vae_pretrainer = None
self.discriminator_pretrainer = None
self.trainer = None
self.decoder_graph = tf.Graph()
self.trained_values = {}
def _build_vae_pretrainer(self, encoder, decoder):
inputs = input_data(shape=self.input_shape, name='input_x')
# Build Network
mean, log_var = encoder(inputs)
encoded = self._encode(mean, log_var)
decoded = decoder(encoded)
# Loss
element_wise_loss = self._get_mean_square(decoded, inputs)
kullback_leibler_divergence = \
self._get_kullback_leibler_divergence(mean, log_var)
pretrain_vae_loss = self.reconstruction_weight_against_detail *\
tf.reduce_mean(element_wise_loss + kullback_leibler_divergence)
# Trainer
pretrain_vae_op = TrainOp(loss=pretrain_vae_loss,
optimizer=self._get_optimizer('vae'),
batch_size=128,
name='VAE_pretrainer')
return Trainer(pretrain_vae_op, tensorboard_dir=PRE_VAE_TENSORBOARD_DIR,
tensorboard_verbose=0,
checkpoint_path=PRE_VAE_CHECKPOINT_PATH,
max_checkpoints=1)
def _build_discriminator_pretrainer(self, decoder, discriminator):
inputs = input_data(shape=self.input_shape, name='input_x')
is_true = input_data(shape=(None, 2), name='is_true')
is_false = input_data(shape=(None, 2), name='is_false')
# Build Network
shape = tf.shape(fc(inputs, self.latent_dim))
random_image = decoder(self._get_z(shape))
prediction_origin, _ = discriminator(inputs)
prediction_random, _ = discriminator(random_image, reuse=True)
# Loss
prediction_all = merge([prediction_origin, prediction_random], 'concat',
axis=0)
y_all = merge([is_true, is_false], 'concat', axis=0)
pretrain_discriminator_loss = crossentropy(prediction_all, y_all)
# Trainer
pretrain_discriminator_op = TrainOp(
loss=pretrain_discriminator_loss,
optimizer=self._get_optimizer('discriminator'), batch_size=128,
trainable_vars=self._get_trainable_variables(discriminator.scope),
name='Discriminator_pretrainer')
return Trainer(pretrain_discriminator_op,
tensorboard_dir=PRE_DIS_TENSORBOARD_DIR,
tensorboard_verbose=0,
checkpoint_path=PRE_DIS_CHECKPOINT_PATH,
max_checkpoints=1)
def _build_trainer(self, encoder, decoder, discriminator):
inputs = input_data(shape=self.input_shape, name='input_x')
is_true = input_data(shape=(None, 2), name='is_true')
is_false = input_data(shape=(None, 2), name='is_false')
# Build Network
mean, log_var = encoder(inputs)
encoded = self._encode(mean, log_var)
decoded = decoder(encoded)
random_image = decoder(self._get_z(tf.shape(mean)), reuse=True)
# Loss
## Encoder
prediction_origin, feature_map_origin = discriminator(inputs)
prediction_decoded, feature_map_decoded = discriminator(decoded, reuse=True)
prediction_random, _ = discriminator(random_image, reuse=True)
## Decoder
feature_wise_loss = \
self._get_mean_square(feature_map_decoded, feature_map_origin)
kullback_leibler_divergence = \
self._get_kullback_leibler_divergence(mean, log_var)
encoder_loss = self.reconstruction_weight_against_detail *\
tf.reduce_mean(feature_wise_loss + kullback_leibler_divergence)
prediction_gan = merge([prediction_decoded, prediction_random],
'concat', axis=0)
y_gan = merge([is_true, is_true], 'concat', axis=0)
gan_generator_loss = crossentropy(prediction_gan, y_gan)
decoder_loss = (encoder_loss + gan_generator_loss) * 0.5
## Discriminator
prediction_fake = merge([prediction_decoded, prediction_random],
'concat', axis=0)
y_fake = merge([is_false, is_false], 'concat', axis=0)
real_loss = crossentropy(prediction_origin, is_true)
fake_loss = crossentropy(prediction_fake, y_fake)
discriminator_loss = (real_loss + fake_loss) * 0.5
# Trainer
encoder_op = TrainOp(
loss=encoder_loss,
optimizer=self._get_optimizer('encoder'),
batch_size=64,
trainable_vars=self._get_trainable_variables(encoder.scope),
name='Encoder')
decoder_op = TrainOp(
loss=decoder_loss,
optimizer=self._get_optimizer('decoder'),
batch_size=64,
trainable_vars=self._get_trainable_variables(decoder.scope),
name='Decoder')
discriminator_op = TrainOp(
loss=discriminator_loss,
optimizer=self._get_optimizer('discriminator'),
batch_size=64,
trainable_vars=self._get_trainable_variables(discriminator.scope),
name='Discriminator')
return Trainer([encoder_op, decoder_op, discriminator_op],
checkpoint_path=CHECKPOINT_PATH, max_checkpoints=1)
def _encode(self, mean, log_var):
epsilon = tf.random_normal(tf.shape(mean), name='Epsilon')
return mean + tf.exp(0.5 * log_var) * epsilon
def _get_z(self, shape):
z = tf.random_normal(shape, name='RandomZ')
return reshape(z, (-1, self.latent_dim))
def _get_kullback_leibler_divergence(self, mean, log_var):
square_mean = tf.pow(mean, 2)
variance = tf.exp(log_var)
kullback_leibler_divergence = \
tf.reduce_mean(1 + log_var - square_mean - variance,
reduction_indices=1)
kullback_leibler_divergence = \
- 0.5 * self.kullback_leibler_ratio * kullback_leibler_divergence
return kullback_leibler_divergence
def _get_mean_square(self, prediction, truth):
return tf.reduce_mean(tf.squared_difference(prediction, truth),
reduction_indices=(1, 2, 3))
def _get_optimizer(self, type_str):
if type_str in ['vae', 'encoder', 'decoder']:
learning_rate = self.vae_learning_rate
beta1 = self.vae_beta1
else: # 'discriminator'
learning_rate = self.discriminator_learning_rate
beta1 = self.discriminator_beta1
opt = adam(learning_rate=learning_rate, beta1=beta1)
return opt.get_tensor()
def _get_trainable_variables(self, scope):
return [v for v in tflearn.get_all_trainable_variable()
if scope + '/' in v.name]
def _get_input_tensor_by_name(self, name):
return tf.get_collection(tf.GraphKeys.INPUTS, scope=name)[0]
def train(self, x, n_sample=None, pretrain_vae_epoch=1,
pretrain_discriminator_epoch=1, train_epoch=10):
if n_sample == None:
n_sample = x.shape[0]
is_true = np.tile([0., 1.], [n_sample, 1])
is_false = np.tile([1., 0.], [n_sample, 1])
encoder = Encoder(self.n_first_channel, self.n_layer, self.latent_dim)
decoder = Decoder(self.img_shape, self.n_first_channel, self.n_layer)
discriminator = Discriminator(self.n_first_channel, self.n_layer)
with tf.Graph().as_default():
self.vae_pretrainer = self._build_vae_pretrainer(encoder, decoder)
trainer = self.vae_pretrainer
input_tensor = self._get_input_tensor_by_name('input_x')
feed_dict = {input_tensor:x}
trainer.fit(feed_dict, n_epoch=pretrain_vae_epoch,
snapshot_epoch=True, shuffle_all=True,
run_id='VAE_pretrain')
self.trained_values[encoder.scope] = \
self._get_trained_values(trainer, encoder.scope)
self.trained_values[decoder.scope] = \
self._get_trained_values(trainer, decoder.scope)
with tf.Graph().as_default():
self.discriminator_pretrainer = \
self._build_discriminator_pretrainer(decoder, discriminator)
trainer = self.discriminator_pretrainer
self._assign_values(trainer, decoder.scope)
input_tensor = self._get_input_tensor_by_name('input_x')
true_tensor = self._get_input_tensor_by_name('is_true')
false_tensor = self._get_input_tensor_by_name('is_false')
feed_dict = {input_tensor:x,
true_tensor:is_true,
false_tensor:is_false}
trainer.fit(feed_dict, n_epoch=pretrain_discriminator_epoch,
snapshot_epoch=True, shuffle_all=True,
run_id='Discriminator_pretrain')
self.trained_values[discriminator.scope] = \
self._get_trained_values(trainer, discriminator.scope)
with tf.Graph().as_default():
self.trainer = self._build_trainer(encoder, decoder, discriminator)
trainer = self.trainer
self._assign_values(trainer, encoder.scope)
self._assign_values(trainer, decoder.scope)
self._assign_values(trainer, discriminator.scope)
self._set_decoder(decoder)
input_tensor = self._get_input_tensor_by_name('input_x')
true_tensor = self._get_input_tensor_by_name('is_true')
false_tensor = self._get_input_tensor_by_name('is_false')
feed_dict = {input_tensor:x,
true_tensor:is_true,
false_tensor:is_false}
self.trainer.fit([feed_dict] * 3, n_epoch=train_epoch,
snapshot_step=1000, snapshot_epoch=False,
shuffle_all=True, run_id='VAEGAN',
callbacks=[CustomCallback(self)])
def _get_trained_values(self, trainer, scope):
return {v.name:tflearn.variables.get_value(v, session=trainer.session)
for v in self._get_trainable_variables(scope)}
def _assign_values(self, trainer, scope):
[trainer.session.run(v.assign(self.trained_values[scope][v.name]))
for v in self._get_trainable_variables(scope)]
def _set_decoder(self, decoder):
with self.decoder_graph.as_default():
inputs = input_data(shape=(None, self.latent_dim))
net = decoder(inputs)
self.decoder = DNN(net)
def decode(self, z):
with self.decoder_graph.as_default():
return self.decoder.predict(z)
class Encoder(object):
def __init__(self, n_first_channel, n_layer, latent_dim):
self.n_first_channel = n_first_channel
self.n_layer = n_layer
self.latent_dim = latent_dim
self.scope = 'Encoder'
def __call__(self, x, reuse=False):
net = x
for i in range(self.n_layer):
n_channel = self.n_first_channel * 2 ** i
net = conv(net, n_channel, 4, strides=2, reuse=reuse,
scope='{s}/Conv_{n}'.format(s=self.scope, n=i))
net = bn(net, reuse=reuse,
scope='{s}/BN_{n}'.format(s=self.scope, n=i))
net = relu(net)
# net = softplus(net)
mean = fc(net, self.latent_dim, reuse=reuse,
scope='{s}/Mean'.format(s=self.scope))
log_var = fc(net, self.latent_dim, reuse=reuse,
scope='{s}/LogVariance'.format(s=self.scope))
return mean, log_var
class Decoder(object):
def __init__(self, img_shape, n_first_channel, n_layer):
self.img_size = img_shape[:2]
self.color_channel = img_shape[2]
self.n_first_channel = n_first_channel * 2 ** (n_layer - 1)
self.n_layer = n_layer
self.scope = 'Decoder'
def __call__(self, z, reuse=False):
net = z
feature_height = self.img_size[0] // 2 ** self.n_layer
feature_width = self.img_size[1] // 2 ** self.n_layer
feature_channel = self.n_first_channel
n_units = feature_height * feature_width * feature_channel
net = fc(net, n_units, reuse=reuse, scope='{s}/FC'.format(s=self.scope))
shape = [-1, feature_height, feature_width, feature_channel]
net = reshape(net, shape)
for i in range(self.n_layer):
feature_height *= 2
feature_width *= 2
if i < self.n_layer - 1:
feature_channel //= 2
else:
feature_channel = self.color_channel
net = bn(net, reuse=reuse,
scope='{s}/BN_{n}'.format(s=self.scope, n=i))
net = relu(net)
# net = elu(net)
net = conv_t(net, feature_channel, 4,
[feature_height, feature_width], strides=2,
reuse=reuse,
scope='{s}/ConvT_{n}'.format(s=self.scope, n=i))
net = sigmoid(net)
return net
class Discriminator(object):
def __init__(self, n_first_channel, n_layer):
self.n_first_channel = n_first_channel
self.n_layer = n_layer
self.scope = 'Discriminator'
def __call__(self, x, reuse=False):
net = x
for i in range(self.n_layer):
net = conv(net, self.n_first_channel * 2 ** i, 4, reuse=reuse,
scope='{s}/Conv_{n}'.format(s=self.scope, n=i))
net = max_pool(net, 2)
net = bn(net, reuse=reuse,
scope='{s}/BN_{n}'.format(s=self.scope, n=i))
# net = relu(net)
net = elu(net)
if i == self.n_layer - 1:
feature_reconstruction = net
net = fc(net, 2, reuse=reuse, scope='{s}/FC'.format(s=self.scope))
net = softmax(net)
return net, feature_reconstruction
class CustomCallback(tflearn.callbacks.Callback):
def __init__(self, model, n_side=10):
self.model = model
self.n_side = n_side
self.sample_z = np.random.normal(size=(n_side ** 2, model.latent_dim))
def _save(self, name, z):
model = self.model
n_side = self.n_side
img_height = model.img_shape[0]
img_width = model.img_shape[1]
img_channel = model.img_shape[2]
image = np.ndarray(shape=(n_side * img_height,
n_side * img_width,
img_channel),
dtype=np.float32)
model.trained_values = {
scope:model._get_trained_values(model.trainer, scope)
for scope in model.trained_values}
with model.decoder_graph.as_default():
[model._assign_values(model.decoder, scope)
for scope in model.trained_values]
decoded = model.decode(z)
for y in range(n_side):
for x in range(n_side):
image[y * img_height : (y + 1) * img_height,
x * img_width : (x + 1) * img_width,
:] = decoded[x + y * n_side]
image = np.clip(image, 0, 1)
image *= 255
io.imsave(name, image.astype(np.uint8))
def on_batch_end(self, training_state, snapshot=False):
if snapshot:
step = training_state.step
file_name = '{path}image-{step}.png'.format(path=CHECKPOINT_PATH,
step=step)
self._save(file_name, self.sample_z)
def on_train_end(self, training_state):
latent_dim = self.model.latent_dim
sample_z = np.ndarray(shape=(self.n_side ** 2, latent_dim),
dtype=np.float32)
for row in range(self.n_side):
start = np.random.normal(size=latent_dim)
stop = np.random.normal(size=latent_dim)
z_rows = np.array([np.linspace(start[i], stop[i], num=self.n_side)
for i in range(latent_dim)]).T
sample_z[row * self.n_side : (row + 1) * self.n_side, :] = z_rows
file_name = '{path}image-final.png'.format(path=CHECKPOINT_PATH)
self._save(file_name, sample_z)
(X, Y), (testX, testY) = tflearn.datasets.cifar10.load_data()
X = np.concatenate((X, testX), axis=0)
Y = np.concatenate((Y, testY), axis=0)
X = X[Y == 1]
img_shape = X.shape[1:]
vaegan = VAEGAN(img_shape=img_shape, n_first_channel=64, n_layer=4,
latent_dim=32, kullback_leibler_ratio=0.01,
reconstruction_weight_against_detail=50.0)
vaegan.train(X, pretrain_vae_epoch=1, pretrain_discriminator_epoch=10,
train_epoch=100)