https://github.com/tolstikhin/adagan
Tip revision: 746bd8e6a5277a3a95463a66f4e631e1b48fad48 authored by itolstikhin on 13 December 2017, 10:56:01 UTC
Fixed plots
Fixed plots
Tip revision: 746bd8e
vae.py
# Copyright 2017 Max Planck Society
# Distributed under the BSD-3 Software license,
# (See accompanying file ./LICENSE.txt or copy at
# https://opensource.org/licenses/BSD-3-Clause)
"""This class implements VAE training.
"""
import os
import logging
import tensorflow as tf
import utils
from utils import ProgressBar
from utils import TQDM
import numpy as np
import ops
from metrics import Metrics
class Vae(object):
"""A base class for running individual VAEs.
"""
def __init__(self, opts, data, weights):
# Create a new session with session.graph = default graph
self._session = tf.Session()
self._trained = False
self._data = data
self._data_weights = np.copy(weights)
# Latent noise sampled ones to apply decoder while training
self._noise_for_plots = utils.generate_noise(opts, 500)
# Placeholders
self._real_points_ph = None
self._noise_ph = None
# Main operations
# FIX
self._loss = None
self._loss_reconstruct = None
self._loss_kl = None
self._generated = None
self._reconstruct_x = None
# Optimizers
self.optim = None
with self._session.as_default(), self._session.graph.as_default():
logging.error('Building the graph...')
self._build_model_internal(opts)
# Make sure AdamOptimizer, if used in the Graph, is defined before
# calling global_variables_initializer().
init = tf.global_variables_initializer()
self._session.run(init)
def __enter__(self):
return self
def __exit__(self, exc_type, exc_value, traceback):
# Cleaning the whole default Graph
logging.error('Cleaning the graph...')
tf.reset_default_graph()
logging.error('Closing the session...')
# Finishing the session
self._session.close()
def train(self, opts):
"""Train a VAE model.
"""
with self._session.as_default(), self._session.graph.as_default():
self._train_internal(opts)
self._trained = True
def sample(self, opts, num=100):
"""Sample points from the trained VAE model.
"""
assert self._trained, 'Can not sample from the un-trained VAE'
with self._session.as_default(), self._session.graph.as_default():
return self._sample_internal(opts, num)
def train_mixture_discriminator(self, opts, fake_images):
"""Train classifier separating true data from points in fake_images.
Return:
prob_real: probabilities of the points from training data being the
real points according to the trained mixture classifier.
Numpy vector of shape (self._data.num_points,)
prob_fake: probabilities of the points from fake_images being the
real points according to the trained mixture classifier.
Numpy vector of shape (len(fake_images),)
"""
with self._session.as_default(), self._session.graph.as_default():
return self._train_mixture_discriminator_internal(opts, fake_images)
def _run_batch(self, opts, operation, placeholder, feed,
placeholder2=None, feed2=None):
"""Wrapper around session.run to process huge data.
It is asumed that (a) first dimension of placeholder enumerates
separate points, and (b) that operation is independently applied
to every point, i.e. we can split it point-wisely and then merge
the results. The second placeholder is meant either for is_train
flag for batch-norm or probabilities of dropout.
TODO: write util function which will be called both from this method
and MNIST classification evaluation as well.
"""
assert len(feed.shape) > 0, 'Empry feed.'
num_points = feed.shape[0]
batch_size = opts['tf_run_batch_size']
batches_num = int(np.ceil((num_points + 0.) / batch_size))
result = []
for idx in xrange(batches_num):
if idx == batches_num - 1:
if feed2 is None:
res = self._session.run(
operation,
feed_dict={placeholder: feed[idx * batch_size:]})
else:
res = self._session.run(
operation,
feed_dict={placeholder: feed[idx * batch_size:],
placeholder2: feed2})
else:
if feed2 is None:
res = self._session.run(
operation,
feed_dict={placeholder: feed[idx * batch_size:
(idx + 1) * batch_size]})
else:
res = self._session.run(
operation,
feed_dict={placeholder: feed[idx * batch_size:
(idx + 1) * batch_size],
placeholder2: feed2})
if len(res.shape) == 1:
# convert (n,) vector to (n,1) array
res = np.reshape(res, [-1, 1])
result.append(res)
result = np.vstack(result)
assert len(result) == num_points
return result
def _build_model_internal(self, opts):
"""Build a TensorFlow graph with all the necessary ops.
"""
assert False, 'VAE base class has no build_model method defined.'
def _train_internal(self, opts):
assert False, 'VAE base class has no train method defined.'
def _sample_internal(self, opts, num):
assert False, 'VAE base class has no sample method defined.'
def _train_mixture_discriminator_internal(self, opts, fake_images):
assert False, 'VAE base class has no mixture discriminator method defined.'
class ImageVae(Vae):
"""A simple VAE implementation, suitable for pictures.
"""
def __init__(self, opts, data, weights):
# One more placeholder for batch norm
self._is_training_ph = None
Vae.__init__(self, opts, data, weights)
def generator(self, opts, noise, is_training, reuse=False, return_logits=False):
"""Generator function, suitable for simple picture experiments.
Args:
noise: [num_points, dim] array, where dim is dimensionality of the
latent noise space.
is_training: bool, defines whether to use batch_norm in the train
or test mode.
return_logits: bool, if true returns the "logits" instead of being
normalized (by tanh or sigmoid depending on "input_normalize_sym".
Returns:
[num_points, dim1, dim2, dim3] array, where the first coordinate
indexes the points, which all are of the shape (dim1, dim2, dim3).
"""
output_shape = self._data.data_shape # (dim1, dim2, dim3)
# Computing the number of noise vectors on-the-go
dim1 = tf.shape(noise)[0]
num_filters = opts['g_num_filters']
num_layers = opts['g_num_layers']
with tf.variable_scope("GENERATOR", reuse=reuse):
height = output_shape[0] / 2**(num_layers - 1)
width = output_shape[1] / 2**(num_layers - 1)
h0 = ops.linear(opts, noise, num_filters * height * width,
scope='h0_lin')
h0 = tf.reshape(h0, [-1, height, width, num_filters])
h0 = tf.nn.relu(h0)
layer_x = h0
for i in xrange(num_layers-1):
scale = 2**(i+1)
_out_shape = [dim1, height * scale, width * scale, num_filters / scale]
layer_x = ops.deconv2d(opts, layer_x, _out_shape, scope='h%d_deconv' % i)
if opts['batch_norm']:
layer_x = ops.batch_norm(opts, layer_x, is_training, reuse, scope='bn%d' % i)
layer_x = tf.nn.relu(layer_x)
if opts['dropout']:
_keep_prob = tf.minimum(
1., 0.9 - (0.9 - keep_prob) * float(i + 1) / (num_layers - 1))
layer_x = tf.nn.dropout(layer_x, _keep_prob)
# # h0 = ops.lrelu(h0)
# _out_shape = [dim1, height * 2, width * 2, num_filters / 2]
# # for 28 x 28 does 7 x 7 --> 14 x 14
# h1 = ops.deconv2d(opts, h0, _out_shape, scope='h1_deconv')
# h1 = ops.batch_norm(opts, h1, is_training, reuse, scope='bn_layer2')
# h1 = tf.nn.relu(h1)
# # h1 = ops.lrelu(h1)
# _out_shape = [dim1, height * 4, width * 4, num_filters / 4]
# # for 28 x 28 does 14 x 14 --> 28 x 28
# h2 = ops.deconv2d(opts, h1, _out_shape, scope='h2_deconv')
# h2 = ops.batch_norm(opts, h2, is_training, reuse, scope='bn_layer3')
# h2 = tf.nn.relu(h2)
# # h2 = ops.lrelu(h2)
_out_shape = [dim1] + list(output_shape)
# data_shape[0] x data_shape[1] x ? -> data_shape
h3 = ops.deconv2d(opts, layer_x, _out_shape,
d_h=1, d_w=1, scope='hlast_deconv')
# h3 = ops.batch_norm(opts, h3, is_training, reuse, scope='bn_layer4')
if return_logits:
return h3
if opts['input_normalize_sym']:
return tf.nn.tanh(h3)
else:
return tf.nn.sigmoid(h3)
def discriminator(self, opts, input_, is_training,
prefix='DISCRIMINATOR', reuse=False):
"""Encoder function, suitable for simple toy experiments.
"""
num_filters = opts['d_num_filters']
with tf.variable_scope(prefix, reuse=reuse):
h0 = ops.conv2d(opts, input_, num_filters / 8, scope='h0_conv')
h0 = ops.batch_norm(opts, h0, is_training, reuse, scope='bn_layer1')
h0 = tf.nn.relu(h0)
h1 = ops.conv2d(opts, h0, num_filters / 4, scope='h1_conv')
h1 = ops.batch_norm(opts, h1, is_training, reuse, scope='bn_layer2')
h1 = tf.nn.relu(h1)
h2 = ops.conv2d(opts, h1, num_filters / 2, scope='h2_conv')
h2 = ops.batch_norm(opts, h2, is_training, reuse, scope='bn_layer3')
h2 = tf.nn.relu(h2)
h3 = ops.conv2d(opts, h2, num_filters, scope='h3_conv')
h3 = ops.batch_norm(opts, h3, is_training, reuse, scope='bn_layer4')
h3 = tf.nn.relu(h3)
# Already has NaNs!!
latent_mean = ops.linear(opts, h3, opts['latent_space_dim'], scope='h3_lin')
log_latent_sigmas = ops.linear(opts, h3, opts['latent_space_dim'], scope='h3_lin_sigma')
return latent_mean, log_latent_sigmas
def _build_model_internal(self, opts):
"""Build the Graph corresponding to VAE implementation.
"""
data_shape = self._data.data_shape
# Placeholders
real_points_ph = tf.placeholder(
tf.float32, [None] + list(data_shape), name='real_points_ph')
noise_ph = tf.placeholder(
tf.float32, [None] + [opts['latent_space_dim']], name='noise_ph')
is_training_ph = tf.placeholder(tf.bool, name='is_train_ph')
lr_decay_ph = tf.placeholder(tf.float32)
# Operations
latent_x_mean, log_latent_sigmas = self.discriminator(
opts, real_points_ph, is_training_ph)
scaled_noise = tf.multiply(
tf.sqrt(1e-6 + tf.exp(log_latent_sigmas)), noise_ph)
loss_kl = 0.5 * tf.reduce_sum(
tf.exp(log_latent_sigmas) +
tf.square(latent_x_mean) -
log_latent_sigmas, axis=1)
if opts['recon_loss'] == 'l2sq':
reconstruct_x = self.generator(opts, latent_x_mean + scaled_noise,
is_training_ph)
loss_reconstruct = tf.reduce_sum(
tf.square(real_points_ph - reconstruct_x), axis=[1,2,3])
loss_reconstruct = loss_reconstruct / 2. / opts['vae_sigma']
elif opts['recon_loss'] == 'cross_entropy':
if opts['input_normalize_sym']:
expected = (real_points_ph + 1.0) / 2.0
else:
expected = real_points_ph
reconstruct_x_logits = self.generator(
opts, latent_x_mean + scaled_noise,
is_training_ph, return_logits=True)
loss_reconstruct = tf.reduce_sum(
tf.nn.sigmoid_cross_entropy_with_logits(
labels=expected, logits=reconstruct_x_logits),
axis=[1,2,3])
else:
raise ValueError("Unknown recon loss value %s" % opts['recon_loss'])
dec_enc_x = self.generator(opts, latent_x_mean,
is_training=False, reuse=True)
loss_reconstruct = tf.reduce_mean(loss_reconstruct)
loss_kl = tf.reduce_mean(loss_kl)
loss = loss_kl + loss_reconstruct
# loss = tf.Print(loss, [loss, loss_kl, loss_reconstruct], 'Loss, KL, reconstruct')
optim = ops.optimizer(opts, decay=lr_decay_ph).minimize(loss)
generated_images = self.generator(opts, noise_ph,
is_training_ph, reuse=True)
self._real_points_ph = real_points_ph
self._noise_ph = noise_ph
self._is_training_ph = is_training_ph
self._optim = optim
self._loss = loss
self._loss_reconstruct = loss_reconstruct
self._lr_decay_ph = lr_decay_ph
self._loss_kl = loss_kl
self._generated = generated_images
self._reconstruct_x = dec_enc_x
self._enc_mean = latent_x_mean
self._enc_log_var = log_latent_sigmas
saver = tf.train.Saver(max_to_keep=10)
tf.add_to_collection('real_points_ph', self._real_points_ph)
tf.add_to_collection('noise_ph', self._noise_ph)
tf.add_to_collection('is_training_ph', self._is_training_ph)
tf.add_to_collection('encoder_mean', self._enc_mean)
tf.add_to_collection('encoder_log_sigma', self._enc_log_var)
tf.add_to_collection('decoder', self._generated)
self._saver = saver
logging.error("Building Graph Done.")
def _train_internal(self, opts):
"""Train a VAE model.
"""
batches_num = self._data.num_points / opts['batch_size']
train_size = self._data.num_points
num_plot = 320
sample_prev = np.zeros([num_plot] + list(self._data.data_shape))
l2s = []
counter = 0
decay = 1.
logging.error('Training VAE')
for _epoch in xrange(opts["gan_epoch_num"]):
if opts['decay_schedule'] == "manual":
if _epoch == 30:
decay = decay / 2.
if _epoch == 50:
decay = decay / 5.
if _epoch == 100:
decay = decay / 10.
if _epoch > 0 and _epoch % opts['save_every_epoch'] == 0:
os.path.join(opts['work_dir'], opts['ckpt_dir'])
self._saver.save(self._session,
os.path.join(opts['work_dir'],
opts['ckpt_dir'],
'trained-pot'),
global_step=counter)
for _idx in xrange(batches_num):
# logging.error('Step %d of %d' % (_idx, batches_num ) )
data_ids = np.random.choice(train_size, opts['batch_size'],
replace=False, p=self._data_weights)
batch_images = self._data.data[data_ids].astype(np.float)
batch_noise = utils.generate_noise(opts, opts['batch_size'])
_, loss, loss_kl, loss_reconstruct = self._session.run(
[self._optim, self._loss, self._loss_kl,
self._loss_reconstruct],
feed_dict={self._real_points_ph: batch_images,
self._noise_ph: batch_noise,
self._lr_decay_ph: decay,
self._is_training_ph: True})
counter += 1
if opts['verbose'] and counter % opts['plot_every'] == 0:
debug_str = 'Epoch: %d/%d, batch:%d/%d' % (
_epoch+1, opts['gan_epoch_num'], _idx+1, batches_num)
debug_str += ' [L=%.2g, Recon=%.2g, KLQ=%.2g]' % (
loss, loss_reconstruct, loss_kl)
logging.error(debug_str)
if opts['verbose'] and counter % opts['plot_every'] == 0:
metrics = Metrics()
points_to_plot = self._run_batch(
opts, self._generated, self._noise_ph,
self._noise_for_plots[0:num_plot],
self._is_training_ph, False)
l2s.append(np.sum((points_to_plot - sample_prev)**2))
metrics.l2s = l2s[:]
metrics.make_plots(
opts,
counter,
None,
points_to_plot,
prefix='sample_e%04d_mb%05d_' % (_epoch, _idx))
reconstructed = self._session.run(
self._reconstruct_x,
feed_dict={self._real_points_ph: batch_images,
self._is_training_ph: False})
metrics.l2s = None
metrics.make_plots(
opts,
counter,
None,
reconstructed,
prefix='reconstr_e%04d_mb%05d_' % (_epoch, _idx))
if opts['early_stop'] > 0 and counter > opts['early_stop']:
break
if _epoch > 0:
os.path.join(opts['work_dir'], opts['ckpt_dir'])
self._saver.save(self._session,
os.path.join(opts['work_dir'],
opts['ckpt_dir'],
'trained-pot-final'),
global_step=counter)
def _sample_internal(self, opts, num):
"""Sample from the trained GAN model.
"""
noise = utils.generate_noise(opts, num)
sample = self._run_batch(
opts, self._generated, self._noise_ph, noise,
self._is_training_ph, False)
return sample
