def _callback(self, sess, info):
if info.epoch % FLAGS.lag != 0:
return
x_real = self.X_test
x_gen = self.generate1(sess, {self.batch_size_ph: FLAGS.mbs})
if FLAGS.arch == 'ae':
x_real = self.ae.decode(x_real, sess)
x_gen = self.ae.decode(x_gen, sess)
name = '{}/small_images_{}.jpg'.format(self.run_name, info.epoch)
utils.save_image_collections(x_gen[np.random.permutation(
FLAGS.mbs)[:50]],
name,
scale_each=True,
shape=(5, 10))
print(
'Epoch {} (total {:.1f}, dist {:.1f}, match {:.1f}, sgd {:.1f} s): loss = {}'
.format(info.epoch, info.time, info.time_gen, info.time_align,
info.time_opt, info.loss))
if FLAGS.arch == 'adv':
print('AE loss = {}, GP = {}'.format(info.reg, info.gp))
if FLAGS.dataset == 'cifar' and (info.epoch <= 50
or info.epoch % 100 == 0):
images = []
for i in range(600 if info.epoch % 100 == 0 else 10):
images.append(self.generate1(sess, {self.batch_size_ph: 100}))
images = np.concatenate(images, axis=0)
images = list((images * 128 + 128).astype(np.int32))
print('Inception score = {}'.format(get_inception_score(images)))
def train_VAE():
global x_train, t_train
# SET UP VAE COMPUTATIONAL GRAPH
optimizer, lower_bound, x_gen, x_input, t_input, n, n_particles = build_VAE(
)
# Define training/evaluation parameters
epochs = 1000
batch_size = 128
iters = x_train.shape[0] // batch_size
save_freq = 10
test_freq = 10
test_batch_size = 400
test_iters = x_test.shape[0] // test_batch_size
result_path = "results/cvae"
# Run the inference
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
for epoch in range(1, epochs + 1):
x_train, t_train = shuffle(x_train, t_train)
lbs = []
for i in range(iters):
x_batch = x_train[i * batch_size:(i + 1) * batch_size]
t_batch = t_train[i * batch_size:(i + 1) * batch_size]
_, lb = sess.run(
[optimizer, lower_bound],
feed_dict={
x_input: x_batch,
t_input: t_batch,
n_particles: 1,
n: batch_size
})
lbs.append(lb)
print("Epoch {}: Lower bound = {}".format(epoch, np.mean(lbs)))
if epoch % save_freq == 0:
t_label = np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
t_label = np.broadcast_to(t_label, (100, 10)).copy()
# Generate images for each class
for j in range(10):
t_label[:, j] = 1
images = sess.run(x_gen,
feed_dict={
t_input: t_label,
n: 100,
n_particles: 1
})
name = os.path.join(result_path + "/label_{}".format(j),
"cvae.epoch.{}.png".format(epoch))
save_image_collections(images, name)
t_label[:, j] = 0
def train(self, sess):
learning_rate_ph = self.learning_rate_ph
optimizer = self.optimizer
train_ae = self.train_ae
saver = self.saver
# Train autoencoder
ckpt_file = tf.train.latest_checkpoint(self.ae_name)
if ckpt_file is not None:
print('Restoring autoencoder...')
saver.restore(sess, ckpt_file)
else:
print('Training autoencoder...')
for epoch in range(1, FLAGS.ae_epoches + 1):
time_epoch = -time.time()
np.random.shuffle(self.X)
losses = []
for t in range(self.X.shape[0] // FLAGS.ae_bs):
x_batch = self.X[t * FLAGS.ae_bs:(t + 1) * FLAGS.ae_bs]
_, l = sess.run(
[train_ae, self.loss],
feed_dict={
self.x:
x_batch,
learning_rate_ph:
FLAGS.ae_lr0 * FLAGS.ae_t0 / (FLAGS.ae_t0 + epoch),
self.is_training:
True
})
losses.append(l)
time_epoch += time.time()
avg_loss = np.mean(losses)
print('Epoch {} ({:.1f})s: Reconstruction loss = {}'.format(
epoch, time_epoch, avg_loss))
save_path = os.path.join(self.ae_name, 'ae.ckpt')
saver.save(sess, save_path)
print('Visualizing autoencoder')
# Visualize AE filters
f = sess.run(self.filters)
f = np.transpose(f, [3, 0, 1, 2])
name = '{}/filter.png'.format(self.ae_name)
utils.save_image_collections(f, name, scale_each=True, shape=(16, 16))
def main():
if args.dataset == 'standard':
X = np.load(
'/home/danyang/mfs/data/hccr/image_1000x20x64x64_stand.npy')
elif args.dataset == 'casia-offline':
X = np.load(
'/home/danyang/mfs/data/hccr/image_1000x300x64x64_casia-offline.npy'
)
elif args.dataset == 'casia-online':
X = np.load(
'/home/danyang/mfs/data/hccr/image_1000x300x64x64_casia-online.npy'
)
else:
print('Unknown Dataset!')
os._exit(-1)
train_x = X[:int(train_ratio * n_y), :int(train_ratio * n_font), :, :]
test_x_font = X[:int(train_ratio * n_y),
int(train_ratio * n_font):n_font, :, :]
test_x_char = X[int(train_ratio * n_y):n_y, :int(train_ratio *
n_font), :, :]
test_x = X[int(train_ratio * n_y):n_y,
int(train_ratio * n_font):n_font, :, :]
epochs = args.epoch
train_batch_size = args.batch_size * FLAGS.num_gpus
learning_rate = args.lr
anneal_lr_freq = 200
anneal_lr_rate = 0.75
result_path = args.result_path
train_iters = min(train_x.shape[0] * train_x.shape[1],
10000) // train_batch_size
is_training = tf.placeholder(tf.bool, shape=[], name='is_training')
x = tf.placeholder(tf.int32, shape=[None, n_x], name='x')
font_source = tf.placeholder(tf.int32,
shape=[None, n_x],
name='font_source')
char_source = tf.placeholder(tf.int32,
shape=[None, n_x],
name='char_source')
pairwise_alpha = tf.placeholder(tf.float32,
shape=[],
name='pairwise_alpha')
learning_rate_ph = tf.placeholder(tf.float32, shape=[], name='lr')
optimizer = tf.train.AdamOptimizer(learning_rate_ph, beta1=0.5)
def build_tower_graph(id_):
tower_x = x[id_ * tf.shape(x)[0] // FLAGS.num_gpus:(id_ + 1) *
tf.shape(x)[0] // FLAGS.num_gpus]
tower_font_source = font_source[id_ * tf.shape(font_source)[0] //
FLAGS.num_gpus:(id_ + 1) *
tf.shape(font_source)[0] //
FLAGS.num_gpus]
tower_char_source = char_source[id_ * tf.shape(char_source)[0] //
FLAGS.num_gpus:(id_ + 1) *
tf.shape(char_source)[0] //
FLAGS.num_gpus]
n = tf.shape(tower_x)[0]
x_obs = tf.tile(tf.expand_dims(tower_x, 0), [1, 1, 1])
def log_joint(observed):
decoder, _, = VLAE(observed, n, is_training)
log_pz_char, log_pz_font, log_px_z = decoder.local_log_prob(
['z_char', 'z_font', 'x'])
return log_pz_char + log_pz_font + log_px_z
encoder, _, _ = q_net(None, tower_x, is_training)
qz_samples_font, log_qz_font = encoder.query('z_font',
outputs=True,
local_log_prob=True)
qz_samples_char, log_qz_char = encoder.query('z_char',
outputs=True,
local_log_prob=True)
encoder, _, _ = q_net(None, tower_font_source, is_training)
qz_samples_font_source, log_qz_font_source = encoder.query(
'z_font', outputs=True, local_log_prob=True)
encoder, _, _ = q_net(None, tower_char_source, is_training)
qz_samples_char_source, log_qz_char_source = encoder.query(
'z_char', outputs=True, local_log_prob=True)
lower_bound = tf.reduce_mean(
zs.iwae(log_joint, {'x': x_obs}, {
'z_font': [qz_samples_font, log_qz_font],
'z_char': [qz_samples_char, log_qz_char]
},
axis=0))
lower_bound_pairwise = pairwise_alpha * tf.reduce_mean(
zs.iwae(log_joint, {'x': x_obs}, {
'z_font': [qz_samples_font_source, log_qz_font_source],
'z_char': [qz_samples_char_source, log_qz_char_source]
},
axis=0))
grads = optimizer.compute_gradients(-lower_bound -
lower_bound_pairwise)
return grads, [lower_bound, lower_bound_pairwise]
tower_losses = []
tower_grads = []
for i in range(FLAGS.num_gpus):
with tf.device('/gpu:%d' % i):
with tf.name_scope('tower_%d' % i):
grads, losses = build_tower_graph(i)
tower_losses.append(losses)
tower_grads.append(grads)
lower_bound, lower_bound_pairwise = multi_gpu.average_losses(tower_losses)
grads = multi_gpu.average_gradients(tower_grads)
infer = optimizer.apply_gradients(grads)
params = tf.trainable_variables()
for i in params:
print(i.name, i.get_shape())
saver = tf.train.Saver(max_to_keep=10,
var_list=tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
scope='encoder') + \
tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
scope='decoder'))
_, z_font, _ = q_net(None, font_source, is_training)
_, _, z_char = q_net(None, char_source, is_training)
_, x_gen = VLAE({
'z_font': z_font,
'z_char': z_char
},
tf.shape(char_source)[0], is_training)
x_gen = tf.reshape(tf.sigmoid(x_gen), [-1, n_xl, n_xl, 1])
with multi_gpu.create_session() as sess:
sess.run(tf.global_variables_initializer())
ckpt_file = tf.train.latest_checkpoint(result_path)
begin_epoch = 1
if ckpt_file is not None:
print('Restoring model from {}...'.format(ckpt_file))
begin_epoch = int(ckpt_file.split('.')[-2]) + 1
saver.restore(sess, ckpt_file)
for epoch in range(begin_epoch, epochs + 1):
if epoch % anneal_lr_freq == 0:
learning_rate *= anneal_lr_rate
time_train = -time.time()
lower_bounds, lower_bounds_pairwise = [], []
x_train = train_x.reshape(-1, n_x)
np.random.shuffle(x_train)
x_train = x_train[:min(train_x.shape[0] * train_x.shape[1], 10000)]
if args.pairwise:
x_font_train = np.tile(
np.expand_dims(
np.array([
X[np.random.randint(0, train_x.shape[0] - 1),
i, :, :] for i in range(train_x.shape[1])
]), 0), (train_x.shape[0], 1, 1, 1))
x_char_train = np.tile(
np.expand_dims(
np.array([
X[i,
np.random.randint(0, train_x.shape[1] - 1), :, :]
for i in range(train_x.shape[0])
]), 1), (1, train_x.shape[1], 1, 1))
x_pair = np.concatenate(
(train_x.reshape(-1, n_x), x_char_train.reshape(
-1, n_x), x_font_train.reshape(-1, n_x)), 1)
np.random.shuffle(x_pair)
x_train = x_pair[:min(train_x.shape[0] *
train_x.shape[1], 10000)]
np.random.shuffle(x_train)
for i in range(train_iters):
if args.pairwise:
_, lb, lbp = sess.run(
[infer, lower_bound, lower_bound_pairwise],
feed_dict={
x:
x_train[i * train_batch_size:(i + 1) *
train_batch_size, :n_x],
char_source:
x_train[i * train_batch_size:(i + 1) *
train_batch_size, n_x:2 * n_x],
font_source:
x_train[i * train_batch_size:(i + 1) *
train_batch_size, 2 * n_x:],
learning_rate_ph:
learning_rate,
pairwise_alpha:
args.pairwise_alpha,
is_training:
True
})
else:
_, lb, lbp = sess.run(
[infer, lower_bound, lower_bound_pairwise],
feed_dict={
x:
x_train[i * train_batch_size:(i + 1) *
train_batch_size],
char_source:
x_train[i * train_batch_size:(i + 1) *
train_batch_size],
font_source:
x_train[i * train_batch_size:(i + 1) *
train_batch_size],
learning_rate_ph:
learning_rate,
is_training:
True
})
lower_bounds.append(lb)
lower_bounds_pairwise.append(lbp)
print('Epoch={} ({:.3f}s/epoch): '
'Lower Bound = {} Lower Bound Pairwise = {}'.format(
epoch, (time.time() + time_train), np.mean(lower_bounds),
np.mean(lower_bounds_pairwise)))
# train reconstruction
gen_images = sess.run(x_gen,
feed_dict={
char_source:
train_x[:10, :10, :, :].reshape(-1, n_x),
font_source:
train_x[:10, :10, :, :].reshape(-1, n_x),
is_training:
False
})
name = "train_{}/VLAE_hccr.epoch.{}.png".format(n_y, epoch)
name = os.path.join(result_path, name)
utils.save_contrast_image_collections(
train_x[:10, :10, :, :].reshape(-1, n_xl, n_xl, 1),
gen_images,
name,
shape=(10, 20),
scale_each=True)
# new font reconstruction
char_index = np.arange(test_x_font.shape[0])
font_index = np.arange(test_x_font.shape[1])
np.random.shuffle(char_index)
np.random.shuffle(font_index)
gen_images = sess.run(x_gen,
feed_dict={
char_source:
test_x_font[char_index[:10], :, :, :]
[:,
font_index[:10], :, :].reshape(-1, n_x),
font_source:
test_x_font[char_index[:10], :, :, :]
[:,
font_index[:10], :, :].reshape(-1, n_x),
is_training:
False
})
name = "test_font_{}/VLAE_hccr.epoch.{}.png".format(n_y, epoch)
name = os.path.join(result_path, name)
utils.save_contrast_image_collections(test_x_font[
char_index[:10], :, :, :][:, font_index[:10], :, :].reshape(
-1, n_xl, n_xl, 1),
gen_images,
name,
shape=(10, 20),
scale_each=True)
# new char reconstruction
char_index = np.arange(test_x_char.shape[0])
font_index = np.arange(test_x_char.shape[1])
np.random.shuffle(char_index)
np.random.shuffle(font_index)
gen_images = sess.run(x_gen,
feed_dict={
char_source:
test_x_char[char_index[:10], :, :, :]
[:,
font_index[:10], :, :].reshape(-1, n_x),
font_source:
test_x_char[char_index[:10], :, :, :]
[:,
font_index[:10], :, :].reshape(-1, n_x),
is_training:
False
})
name = "test_char_{}/VLAE_hccr.epoch.{}.png".format(n_y, epoch)
name = os.path.join(result_path, name)
utils.save_contrast_image_collections(test_x_char[
char_index[:10], :, :, :][:, font_index[:10], :, :].reshape(
-1, n_xl, n_xl, 1),
gen_images,
name,
shape=(10, 20),
scale_each=True)
# never seen reconstruction
char_index = np.arange(test_x.shape[0])
font_index = np.arange(test_x.shape[1])
np.random.shuffle(char_index)
np.random.shuffle(font_index)
gen_images = sess.run(x_gen,
feed_dict={
char_source:
test_x[char_index[:10], :, :, :]
[:,
font_index[:10], :, :].reshape(-1, n_x),
font_source:
test_x[char_index[:10], :, :, :]
[:,
font_index[:10], :, :].reshape(-1, n_x),
is_training:
False
})
name = "test_{}/VLAE_hccr.epoch.{}.png".format(n_y, epoch)
name = os.path.join(result_path, name)
utils.save_contrast_image_collections(test_x[
char_index[:10], :, :, :][:, font_index[:10], :, :].reshape(
-1, n_xl, n_xl, 1),
gen_images,
name,
shape=(10, 20),
scale_each=True)
# one shot font generation
font_index = np.arange(test_x_font.shape[1])
np.random.shuffle(font_index)
test_x_font_feed = np.tile(
np.expand_dims(
np.array([
test_x_font[np.random.randint(test_x_font.shape[0] -
1), font_index[i], :, :]
for i in range(10)
]), 0), (10, 1, 1, 1))
gen_images = sess.run(x_gen,
feed_dict={
char_source:
train_x[:10, :10, :, :].reshape(-1, n_x),
font_source:
test_x_font_feed[:10, :10, :, :].reshape(
-1, n_x),
is_training:
False
})
images = np.concatenate(
[test_x_font_feed[0].reshape(-1, n_xl, n_xl, 1), gen_images],
0)
name = "one_shot_font_{}/VLAE_hccr.epoch.{}.png".format(n_y, epoch)
name = os.path.join(result_path, name)
utils.save_image_collections(images,
name,
shape=(11, 10),
scale_each=True)
# one shot char generation
char_index = np.arange(test_x_char.shape[0])
np.random.shuffle(char_index)
test_x_char_feed = np.tile(
np.expand_dims(
np.array([
test_x_char[char_index[i],
np.random.randint(test_x_char.shape[1] -
1), :, :]
for i in range(10)
]), 1), (1, 10, 1, 1))
gen_images = sess.run(x_gen,
feed_dict={
char_source:
test_x_char_feed[:10, :10, :, :].reshape(
-1, n_x),
font_source:
train_x[:10, :10, :, :].reshape(-1, n_x),
is_training:
False
})
name = "one_shot_char_{}/VLAE_hccr.epoch.{}.png".format(n_y, epoch)
name = os.path.join(result_path, name)
images = np.zeros((110, 64, 64, 1))
for i in range(10):
images[i * 11] = np.expand_dims(test_x_char_feed[i, 0, :, :],
2)
images[i * 11 + 1:(i + 1) * 11] = gen_images[i * 10:(i + 1) *
10]
utils.save_image_collections(images,
name,
shape=(10, 11),
scale_each=True)
save_path = "VLAE.epoch.{}.ckpt".format(epoch)
save_path = os.path.join(result_path, save_path)
saver.save(sess, save_path)
time_test = -time.time()
# train gen
t_test = np.reshape(t_test, (-1, n_code))
t_batch = t_test[:gen_size]
gen_images = sess.run(eval_x_gen,
feed_dict={
is_training: False,
code: t_batch
})
name = "gen_{}/iwae_hccr.epoch.{}.iter.{}.png".format(
n_y, epoch, iter)
name = os.path.join(result_path, name)
utils.save_image_collections(
gen_images,
name,
shape=(test_ny, display_each_character),
scale_each=True)
# # train reconstruction
# x_batch = x_train_recon[:recon_size].reshape(-1, n_x)
# x_batch_bin = sess.run(x_bin, feed_dict={x_orig: x_batch})
# t_batch = t_train_recon[:recon_size]
# eval_zs, recon_images = \
# sess.run([eval_z_gen.tensor, eval_x_recon],
# feed_dict={x: x_batch_bin, is_training: False, code: t_batch})
# name = "train_recon_{}/iwae_hccr.epoch.{}.iter.{}.png".format(n_y,
# epoch, iter)
# recon_images = (recon_images > print_threhold).astype(np.float32)
# name = os.path.join(
# result_path, name)
if iter % test_freq == 0:
time_test = -time.time()
t_batch = t_test[:gen_size]
gen_images = sess.run(eval_x_gen,
feed_dict={
is_training: False,
code: t_batch
})
name = "gen_{}/iwae_hccr.epoch.{}.iter.{}.png".format(
n_y, epoch, iter)
name = os.path.join(result_path, name)
utils.save_image_collections(
gen_images,
name,
shape=(test_ny, display_each_character),
scale_each=True)
# train reconstruction
x_batch = x_train_recon[:recon_size].reshape(-1, n_x)
x_batch_bin = sess.run(x_bin, feed_dict={x_orig: x_batch})
t_batch = t_train_recon[:recon_size]
eval_zs, recon_images = \
sess.run([eval_z_gen.tensor, eval_x_recon],
feed_dict={x: x_batch_bin, is_training: False, code: t_batch})
name = "train_recon_{}/iwae_hccr.epoch.{}.iter.{}.png".format(
n_y, epoch, iter)
recon_images = (recon_images > print_threhold).astype(
np.float32)
name = os.path.join(result_path, name)
def main():
seed = FLAGS.seed
result_path = "results/mnist_crowd_{}_{}".format(time.strftime("%Y%m%d_%H%M%S"), seed)
logger = setup_logger('mnist', __file__, result_path)
np.random.seed(seed)
tf.set_random_seed(seed)
# Load MNIST
data_path = os.path.join('data', 'mnist.pkl.gz')
o_train, t_train, o_valid, t_valid, o_test, t_test = \
dataset.load_mnist_realval(data_path, one_hot=False)
o_train = np.vstack([o_train, o_valid])
t_train = np.hstack([t_train, t_valid])
n_train, o_dim = o_train.shape
# indices = np.random.permutation(n_train)
# o_train = o_train[indices]
# t_train = t_train[indices]
o_test = np.random.binomial(1, o_test, size=o_test.shape)
n_test, _ = o_test.shape
# n_class = np.max(t_test) + 1
# Prior parameters
d = 8
K = 50
W = 20
prior_alpha = 1.05
prior_niw_conc = 0.5
prior_tau = 1.
# Variational initialization
alpha = 2.
niw_conc = 1.
random_scale = 3.
tau = 10.
# learning rate
learning_rate = 1e-3
nat_grad_scale = 1e4
# Load annotations
# [i, j, w, L]
annotations = load_annotations(t_train, W, method="real")
n_annotations = annotations.shape[0]
W = len(set(annotations[:, 2]))
# batch_size = 128
# iters = o_train.shape[0] // batch_size
# ann_batch_size = annotations.shape[0] // iters
# print(ann_batch_size)
# exit(0)
# Define training parameters
epochs = 200
batch_size = 128
iters = o_train.shape[0] // batch_size
ann_batch_size = annotations.shape[0] // iters
save_freq = 1
test_freq = 10
test_batch_size = 400
test_iters = o_test.shape[0] // test_batch_size
prior_global_params = get_global_params(
"prior", d, K, W, prior_alpha, prior_niw_conc, prior_tau,
trainable=False)
global_params = get_global_params(
"variational", d, K, W, alpha, niw_conc, tau,
random_scale=random_scale, trainable=True)
# n_particles = tf.placeholder(tf.int32, shape=[], name='n_particles')
o_input = tf.placeholder(tf.float32, shape=[None, o_dim], name='o')
o = tf.to_int32(tf.random_uniform(tf.shape(o_input)) <= o_input)
ann_o_input = tf.placeholder(tf.float32, shape=[None, o_dim], name='ann_o')
ann_o = tf.to_int32(tf.random_uniform(tf.shape(ann_o_input)) <= ann_o_input)
L_ph = tf.sparse_placeholder(tf.float32, shape=[None, None, W])
I_ph = tf.sparse_placeholder(tf.float32, shape=[None, None, W])
lower_bound, global_nat_grads, z_stats, niw_stats, dir_stats = \
variational_message_passing(
prior_global_params, global_params, o, o_dim, d, K, n_train,
n_iters=4)
z_pred = tf.argmax(z_stats, axis=-1)
ann_lower_bound, ann_nat_grads, _, _, _ = variational_message_passing(
prior_global_params, global_params, ann_o, o_dim, d, K, n_train,
L_ph, I_ph, n_annotations, ann_batch_size, n_iters=4)
# ann_lower_bound = tf.constant(0.)
# ann_nat_grads = [tf.zeros_like(param) for param in global_params]
optimizer = tf.train.MomentumOptimizer(learning_rate=learning_rate,
momentum=0.9)
net_vars = (tf.trainable_variables(scope="encoder") +
tf.trainable_variables(scope="decoder"))
net_grads_and_vars = optimizer.compute_gradients(
-0.5 * (lower_bound + ann_lower_bound), var_list=net_vars)
global_nat_grads.extend([0, 0])
nat_grads = [-nat_grad_scale * 0.5 * (g + ann_g)
for g, ann_g in zip(global_nat_grads, ann_nat_grads)]
global_grads_and_vars = list(zip(nat_grads, global_params))
infer_op = optimizer.apply_gradients(net_grads_and_vars +
global_grads_and_vars)
# Generation
# niw_stats: [K, d + d^2 + 2]
gen_mvn_params = niw_stats[:, :-2]
# transparency: [K]
transp = tf.exp(dir_stats) / tf.reduce_max(tf.exp(dir_stats))
# x_samples: [K, d, 10]
x_samples = mvn.sample(gen_mvn_params, d, n_samples=10)
# o_mean: [10, K, o_dim]
_, o_mean = decoder(tf.transpose(x_samples, [2, 0, 1]), o_dim)
# o_gen: [10 * K, 28, 28, 1]
o_gen = tf.reshape(o_mean * transp[:, None], [-1, 28, 28, 1])
def _evaluate(pred_batches, labels):
preds = np.hstack(pred_batches)
truths = labels[:preds.size]
acc, _ = cluster_acc(preds, truths)
nmi = adjusted_mutual_info_score(truths, labels_pred=preds)
return acc, nmi
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
for epoch in range(1, epochs + 1):
time_epoch = -time.time()
indices = np.random.permutation(n_train)
# print(indices[:5])
# exit(0)
o_train_raw = o_train[indices]
t_train_raw = t_train[indices]
lbs, ann_lbs = [], []
t_preds, ann_t_preds = [], []
for t in range(iters):
# Without annotation
o_batch = o_train_raw[t * batch_size:(t + 1) * batch_size]
# With annotation
ann_indices = np.random.randint(0, n_annotations,
size=ann_batch_size)
ann_batch = annotations[ann_indices]
o_indices, orig_to_batch_ind, batch_to_orig_ind, \
sparse_ann_batch, sparse_ann_ind = make_sparse_ann_batch(
ann_batch, W)
ann_o_batch = o_train[o_indices]
_, lb, t_pred, ann_lb = sess.run(
[infer_op, lower_bound, z_pred, ann_lower_bound],
feed_dict={o_input: o_batch,
ann_o_input: ann_o_batch,
L_ph: sparse_ann_batch,
I_ph: sparse_ann_ind})
lbs.append(lb)
t_preds.append(t_pred)
# print("lb: {}".format(lb))
ann_lbs.append(ann_lb)
time_epoch += time.time()
train_acc, train_nmi = _evaluate(t_preds, t_train_raw)
logger.info(
'Epoch {} ({:.1f}s): Lower bound = {}, ann LB = {}, '
'acc = {}, nmi = {}'
.format(epoch, time_epoch, np.mean(lbs), np.mean(ann_lbs),
train_acc, train_nmi))
if epoch % test_freq == 0:
time_test = -time.time()
test_lbs = []
test_t_preds = []
for t in range(test_iters):
test_o_batch = o_test[t * test_batch_size:
(t + 1) * test_batch_size]
test_lb, test_t_pred = sess.run([lower_bound, z_pred],
feed_dict={o: test_o_batch})
test_lbs.append(test_lb)
test_t_preds.append(test_t_pred)
time_test += time.time()
test_acc, test_nmi = _evaluate(test_t_preds, t_test)
logger.info('>>> TEST ({:.1f}s)'.format(time_test))
logger.info('>> Test lower bound = {}, acc = {}, nmi = {}'
.format(np.mean(test_lbs), test_acc, test_nmi))
if epoch == epochs:
with open('results/mnist_bayesSCDC.txt', "a") as myfile:
myfile.write("seed: %d train_acc: %f train_nmi: %f "
"test_acc: %f test_nmi: %f" % (
seed, train_acc, train_nmi, test_acc, test_nmi))
myfile.write('\n')
myfile.close()
if epoch % save_freq == 0:
logger.info('Saving images...')
images = sess.run(o_gen)
name = os.path.join(result_path,
"vae.epoch.{}.png".format(epoch))
save_image_collections(images, name, shape=(10, K))
if args.dataset == 'standard':
data = np.load(
'/home/danyang/mfs/data/hccr/image_1000x200x64x64_stand.npy')
elif args.dataset == 'casia-online':
data = np.load(
'/home/danyang/mfs/data/hccr/image_1000x300x64x64_casia-online.npy')
elif args.dataset == 'casia-offline':
data = np.load(
'/home/danyang/mfs/data/hccr/image_1000x300x64x64_casia-offline.npy')
data = 1.0 - data
print np.max(data), np.min(data)
#np.save('/home/danyang/mfs/data/hccr/image_1000x300x64x64_casia-offline-reverse.npy' , data)
elif args.dataset == 'sogou':
data = np.load(
'/home/danyang/mfs/data/hccr/image_500x1000x64x64_sogou.npy')
else:
print 'dataset error!'
n_y = np.shape(data)[0]
handwritten_index = [3, 9, 10, 20, 194, 195, 196, 197, 198, 193]
standard_index = [0, 1, 5, 13, 18, 31, 39, 44, 182, 168]
sample_num = np.shape(data)[1]
char_num = 100
font_num = 500
display = np.reshape(data[:char_num, :font_num, :, :], (-1, 64, 64, 1))
#np.random.shuffle(display)
utils.save_image_collections(display,
'./result/%s.png' % args.dataset,
shape=(char_num, 300),
scale_each=True)
def save_img(sess, x, z, ckpt_path, epoch):
x_val, z_val = sess.run([x, z])
img_path = os.path.join(ckpt_path, "result.epoch{}.png".format(epoch))
save_image_collections(x_val, img_path)
return z_val
def main():
# Load MNIST
data_path = os.path.join(conf.data_dir, "mnist.pkl.gz")
x_train, t_train, x_valid, t_valid, x_test, t_test = \
dataset.load_mnist_realval(data_path)
x_train = np.vstack([x_train, x_valid])
y_train = np.vstack([t_train, t_valid])
x_test = np.random.binomial(1, x_test, size=x_test.shape)
x_dim = x_train.shape[1]
y_dim = y_train.shape[1]
# Define model parameters
z_dim = 40
# Build the computation graph
n_particles = tf.placeholder(tf.int32, shape=[], name="n_particles")
x_input = tf.placeholder(tf.float32, shape=[None, x_dim], name="x")
x = tf.cast(tf.less(tf.random_uniform(tf.shape(x_input)), x_input),tf.int32)
y = tf.placeholder(tf.float32, shape=[None, 10], name="y")
n = tf.placeholder(tf.int32, shape=[], name="n")
model = build_gen(x_dim, z_dim, y, n, n_particles)
variational = build_q_net(x, z_dim, y, n_particles)
lower_bound = zs.variational.elbo(
model, {"x": x}, variational=variational, axis=0)
cost = tf.reduce_mean(lower_bound.sgvb())
lower_bound = tf.reduce_mean(lower_bound)
# # Importance sampling estimates of marginal log likelihood
is_log_likelihood = tf.reduce_mean(
zs.is_loglikelihood(model, {"x": x}, proposal=variational, axis=0))
optimizer = tf.train.AdamOptimizer(learning_rate=0.001)
infer_op = optimizer.minimize(cost)
# Random generation
y_observe = tf.placeholder(tf.float32, shape=[None, 10], name="y_observe")
x_gen = tf.reshape(model.observe(y=y_observe)["x_mean"], [-1, 28, 28, 1])
# Define training/evaluation parameters
epochs = 3000
batch_size = 128
iters = x_train.shape[0] // batch_size
save_freq = 10
test_freq = 10
test_batch_size = 400
test_iters = x_test.shape[0] // test_batch_size
result_path = "results/vae"
# Run the inference
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
for epoch in range(1, epochs + 1):
time_epoch = -time.time()
np.random.shuffle(x_train)
lbs = []
for t in range(iters):
x_batch = x_train[t * batch_size:(t + 1) * batch_size]
y_batch = y_train[t * batch_size:(t + 1) * batch_size]
_, lb = sess.run([infer_op, lower_bound],
feed_dict={x_input: x_batch,
y: y_batch,
n_particles: 1,
n: batch_size})
lbs.append(lb)
time_epoch += time.time()
print("Epoch {} ({:.1f}s): Lower bound = {}".format(
epoch, time_epoch, np.mean(lbs)))
if epoch % test_freq == 0:
time_test = -time.time()
test_lbs, test_lls = [], []
for t in range(test_iters):
test_x_batch = x_test[t * test_batch_size:(t + 1) * test_batch_size]
test_y_batch = t_test[t * test_batch_size:(t + 1) * test_batch_size]
test_lb = sess.run(lower_bound,
feed_dict={x: test_x_batch,
y: test_y_batch,
n_particles: 1,
n: test_batch_size})
test_ll = sess.run(is_log_likelihood,
feed_dict={x: test_x_batch,
y: test_y_batch,
n_particles: 1000,
n: test_batch_size})
test_lbs.append(test_lb)
test_lls.append(test_ll)
time_test += time.time()
print(">>> TEST ({:.1f}s)".format(time_test))
print(">> Test lower bound = {}".format(np.mean(test_lbs)))
print('>> Test log likelihood (IS) = {}'.format(
np.mean(test_lls)))
if epoch % save_freq == 0:
y_index = np.repeat(np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]), 10)
y_index = np.eye(10)[y_index.reshape(-1)]
images = sess.run(x_gen, feed_dict={y_observe: y_index, y: y_index, n: 100, n_particles: 1})
name = os.path.join(result_path, "c-vae.epoch.{}.png".format(epoch))
save_image_collections(images, name)
batch_size = 200
iters = x_train.shape[0] // batch_size
learning_rate = 0.0002
anneal_lr_freq = 200
anneal_lr_rate = 0.75
test_freq = 10
test_batch_size = 400
test_iters = x_test.shape[0] // test_batch_size
save_freq = 100
result_path = "result/VLAE-mnist"
index = np.arange(60000)
np.random.shuffle(index)
x_dis = x_train[index[:400], :].reshape(400, 28, 28, n_channel)
utils.save_image_collections(x_dis,
os.path.join(result_path, 'data.png'),
scale_each=True)
# Build the computation graph
is_training = tf.placeholder(tf.bool, shape=[], name='is_training')
n_particles = tf.placeholder(tf.int32, shape=[], name='n_particles')
x_orig = tf.placeholder(tf.float32, shape=[None, n_x], name='x')
x_bin = tf.cast(tf.less(tf.random_uniform(tf.shape(x_orig), 0, 1), x_orig),
tf.int32)
x = tf.placeholder(tf.int32, shape=[None, n_x], name='x')
x_obs = tf.tile(tf.expand_dims(x, 0), [n_particles, 1, 1])
n = tf.shape(x)[0]
#gen
gen_z_0 = tf.placeholder(tf.float32, shape=[400, n_z_0], name='gen_z_0')
gen_z_1 = tf.placeholder(tf.float32, shape=[400, n_z_1], name='gen_z_1')
def get_code():
for i in params:
print(i.name, i.get_shape())
var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES,
scope='decoder') + \
tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES,
scope='encoder')
saver = tf.train.Saver(max_to_keep=10, var_list=var_list)
with multi_gpu.create_session() as sess:
sess.run(tf.global_variables_initializer())
ckpt_file = tf.train.latest_checkpoint(result_path)
begin_epoch = 1
if ckpt_file is not None:
print('Restoring model from {}...'.format(ckpt_file))
begin_epoch = int(ckpt_file.split('.')[-4]) + 1
saver.restore(sess, ckpt_file)
for epoch in range(begin_epoch, epoches + 1):
if epoch % anneal_lr_freq == 0:
learning_rate *= anneal_lr_rate
x_train_source, t_train_source = utils.random_select(x_train.reshape(n_y, -1, n_x), t_train)
x_train_source = np.tile(x_train_source,(n_y,1))
t_train_source = np.tile(t_train_source,(n_y,1))
x_train_tmp, t_train_tmp = utils.shuffle(
np.concatenate((x_train.reshape(-1, n_x, 1), x_train_source.reshape(-1, n_x, 1)), axis=2),
np.concatenate((t_train.reshape(-1, n_code, 1), t_train_source.reshape(-1, n_code, 1)), axis=2))
x_train_shuffle = x_train_tmp[:, :, 0].reshape(-1, n_x)
t_train_shuffle = t_train_tmp[:, :, 0].reshape(-1, n_code)
x_train_source = x_train_tmp[:, :, 1].reshape(-1, n_x)
t_train_source = t_train_tmp[:, :, 1].reshape(-1, n_code)
# x_train = np.reshape(x_train, [-1, n_xl, n_xl, 1])
lower_bounds = []
tv_losses = []
time_train = -time.time()
for t in range(train_iters):
iter = t + 1
x_batch = x_train_shuffle[t * train_batch_size:(t + 1) * train_batch_size]
x_batch_bin = sess.run(x_bin, feed_dict={x_orig: x_batch})
t_batch = t_train_shuffle[t * train_batch_size:(t + 1) * train_batch_size]
x_batch_source = x_train_source[t * train_batch_size:(t + 1) * train_batch_size]
x_batch_source_bin = sess.run(x_bin, feed_dict={x_orig:x_batch_source})
t_batch_source = t_train_source[t * train_batch_size:(t + 1) * train_batch_size]
_, lb , tv = sess.run([infer, lower_bound , tv_loss],
feed_dict={x_orig: x_batch, x: x_batch_bin, code: t_batch,
x_source: x_batch_source_bin, code_source: t_batch_source,
learning_rate_ph: learning_rate, is_training: True})
lower_bounds.append(lb)
tv_losses.append(tv)
if iter % print_freq == 0:
print('Epoch={} Iter={} ({:.3f}s/iter): '
'Lower Bound={} , Tv loss={}'.
format(epoch, iter,
(time.time() + time_train) / print_freq,
np.mean(lower_bounds) , np.mean(tv_losses)))
lower_bounds = []
tv_losses = []
if iter % test_freq == 0:
time_test = -time.time()
t_batch = t_test[:gen_size]
gen_images = sess.run(eval_x_gen,
feed_dict={is_training: False,
code: t_batch})
name = "gen_{}/iwae_hccr.epoch.{}.iter.{}.png".format(n_y, epoch, iter)
name = os.path.join(result_path, name)
utils.save_image_collections(gen_images, name, shape=(test_ny, display_each_character),
scale_each=True)
# train reconstruction
x_batch = x_train_recon[:recon_size].reshape(-1, n_x)
x_batch_bin = sess.run(x_bin, feed_dict={x_orig: x_batch})
t_batch = t_train_recon[:recon_size]
eval_zs, recon_images = \
sess.run([eval_z_gen.tensor, eval_x_recon],
feed_dict={x: x_batch_bin, is_training: False, code: t_batch})
name = "train_recon_{}/iwae_hccr.epoch.{}.iter.{}.png".format(n_y,
epoch, iter)
name = os.path.join(
result_path, name)
utils.save_contrast_image_collections(x_batch.reshape(-1, n_xl, n_xl, n_channels), recon_images,
name, shape=(test_ny, display_each_character * 2),
scale_each=True)
# # train interpolation
x_batch_bin = sess.run(x_bin, feed_dict={x_orig: x_train_interp})
t_batch = t_train_interp
eval_zs, _ = \
sess.run([eval_z_gen.tensor, eval_x_recon],
feed_dict={x: x_batch_bin, is_training: False, code: t_batch})
epsilon = np.linspace(0, 1, display_each_character)
eval_zs_interp = np.array(
[eps * eval_zs[0, 2 * i, :] + (1 - eps) * eval_zs[0, 2 * i + 1, :] for i in range(test_ny) for eps
in epsilon]).reshape(1, -1, n_z)
t_batch = np.tile([t_batch[2 * i, :] for i in range(test_ny)], (1, display_each_character)).reshape(-1,
n_code)
recon_images = \
sess.run(eval_x_interp, feed_dict={interp_z: eval_zs_interp, is_training: False, code: t_batch})
name = "interp_{}/iwae_hccr.epoch.{}.iter.{}.png".format(n_y, epoch, iter)
name = os.path.join(result_path, name)
utils.save_image_collections(recon_images, name, shape=(test_ny, display_each_character),
scale_each=True)
# test reconstruction
x_batch = x_test[:recon_size].reshape(-1, n_x)
x_batch_bin = sess.run(x_bin, feed_dict={x_orig: x_batch})
t_batch = t_test[:recon_size]
eval_zs, recon_images = \
sess.run([eval_z_gen.tensor, eval_x_recon],
feed_dict={x: x_batch_bin, is_training: False, code: t_batch})
name = "test_recon_{}/iwae_hccr.epoch.{}.iter.{}.png".format(n_y,
epoch, iter)
name = os.path.join(
result_path, name)
utils.save_contrast_image_collections(x_batch.reshape(-1, n_xl, n_xl, n_channels), recon_images,
name, shape=(test_ny, display_each_character * 2),
scale_each=True)
# one-shot generation
x_batch = x_oneshot_test.reshape(-1, n_x) # display_number*nxl*nxl*nchannel
x_batch_bin = sess.run(x_bin, feed_dict={x_orig: x_batch})
t_batch = t_oneshot_test
display_x_oneshot = np.zeros((display_each_character, test_ny+ 1, n_xl, n_xl, n_channels))
eval_zs_oneshot = sess.run(eval_z_oneshot.tensor,
feed_dict={x: x_batch_bin, is_training: False, code: t_batch})
# print (np.shape(eval_zs_oneshot)) #test_ny*nz
for i in range(display_each_character):
display_x_oneshot[i, 0, :, :, :] = x_batch[i, :].reshape(-1, n_xl, n_xl, n_channels)
tmp_z = np.zeros((1, test_ny, n_z))
for j in range(test_ny):
# print (np.shape(tmp_z) ,np.shape(eval_zs_oneshot))
tmp_z[0, j, :] = eval_zs_oneshot[0, i, :]
# _, eval_x_oneshot = decoder({'z': oneshot_z}, tf_ny, code, is_training)
#print tmp_z.shape , t_oneshot_gen_test.shape
tmp_x = sess.run(eval_x_oneshot,
feed_dict={oneshot_z: tmp_z, tf_ny: test_ny, code: t_oneshot_gen_test,
is_training: False})
# print (np.shape(tmp_x))
display_x_oneshot[i, 1:, :, :, :] = tmp_x
display_x_oneshot = np.reshape(display_x_oneshot, (-1, n_xl, n_xl, n_channels))
##TODO
display_x_oneshot = (display_x_oneshot > 0.5).astype(np.float32)
name = "oneshot_{}/iwae_hccr.epoch.{}.iter.{}.png".format(n_y,
epoch, iter)
name = os.path.join(
result_path, name)
utils.save_image_collections(display_x_oneshot,
name, shape=(display_each_character, test_ny + 1),
scale_each=True)
# disentangle
t_batch = t_test[:recon_size]
z_each = np.random.normal(size=(display_each_character, n_z))
# print (z_each.shape)
z_batch = np.zeros((test_ny, display_each_character, n_z))
# print (z_batch.shape)
for i in range(test_ny):
z_batch[i, :, :] = z_each
z_batch = np.reshape(z_batch, (-1, n_z))
eval_disentange_x = \
sess.run(disentangle_x,
feed_dict={disentange_z: z_batch, is_training: False, code: t_batch})
name = "disentangle_{}/iwae_hccr.epoch.{}.iter.{}.png".format(n_y,
epoch, iter)
name = os.path.join(
result_path, name)
utils.save_image_collections(eval_disentange_x,
name, shape=(test_ny, display_each_character),
scale_each=True)
time_test += time.time()
if iter % save_freq == 0:
save_path = "iwae.epoch.{}.iter.{}.ckpt".format(epoch, iter)
save_path = os.path.join(result_path, save_path)
saver.save(sess, save_path)
if iter % print_freq == 0:
time_train = -time.time()
def train_vae(args):
# Load MNIST
data_path = os.path.join(args.data_dir, "mnist.pkl.gz")
x_train, y_train, x_valid, y_valid, x_test, y_test = dataset.load_mnist_realval(data_path)
x_train = np.random.binomial(1, x_train, size=x_train.shape)
x_dim = x_train.shape[1]
y_dim = y_train.shape[1]
# Define model parameters
z_dim = args.z_dim
# Build the computation graph
x = tf.placeholder(tf.float32, shape=[None, x_dim], name="x")
y = tf.placeholder(tf.float32, shape=[None, y_dim], name="y")
n = tf.placeholder(tf.int32, shape=[], name="n")
# Get the models
model = build_gen(y, x_dim, z_dim, n)
variational = build_q_net(x, y, z_dim)
# Calculate ELBO
lower_bound = zs.variational.elbo(model, {"x": x }, variational=variational)
cost = tf.reduce_mean(lower_bound.sgvb())
lower_bound = tf.reduce_mean(lower_bound)
optimizer = tf.train.AdamOptimizer(learning_rate=args.lr)
infer_op = optimizer.minimize(cost)
# Random generation
x_gen = tf.reshape(model.observe()["x_mean"], [-1, 28, 28, 1])
# Compute class labels
labels = []
for c in range(10):
l = np.zeros((100, 10))
l[:,c] = 1
labels.append(l)
epochs = args.epochs
batch_size = args.batch_size
iters = x_train.shape[0] // batch_size
saver = tf.train.Saver(max_to_keep=10)
save_model_freq = min(100, args.epochs)
# Run the Inference
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
ckpt_file = tf.train.latest_checkpoint(args.checkpoints_path)
begin_epoch = 1
if(ckpt_file is not None):
print('Restoring model from {}...'.format(ckpt_file))
begin_epoch = int(ckpt_file.split('.')[-2]) + 1
saver.restore(sess, ckpt_file)
for epoch in range(1, epochs+1):
time_epoch = -time.time()
lbs = []
for t in range(iters):
x_batch = x_train[t*batch_size:(t+1)*batch_size]
y_batch = y_train[t*batch_size:(t+1)*batch_size]
_, lb = sess.run(
[infer_op, lower_bound],
feed_dict={
x: x_batch,
y: y_batch,
n: batch_size
}
)
lbs.append(lb)
time_epoch += time.time()
print("Epoch {} ({:.1f}s): Lower bound = {}".format(epoch, time_epoch, np.mean(lbs)))
if(epoch % args.save_model_freq == 0):
save_path = os.path.join(args.checkpoints_path, "vae.epoch.{}.ckpt".format(epoch))
if not os.path.exists(os.path.dirname(save_path)):
os.makedirs(os.path.dirname(save_path))
saver.save(sess, save_path)
if epoch % args.save_img_freq == 0:
for c in range(10):
images = sess.run(x_gen, feed_dict={y: labels[c], n: 100 })
name = os.path.join(args.results_path, str(epoch).zfill(3), "{}.png".format(c))
utils.save_image_collections(images, name)
time_test = -time.time()
tt_lbs = []
tt_preds = []
for tt in range(test_iters):
test_x_batch = x_test_bin[tt * tt_bs:(tt + 1) * tt_bs]
tt_pred, tt_lb = sess.run([z_pred, lower_bound],
feed_dict={x: test_x_batch})
tt_preds.append(tt_pred)
tt_lbs.append(tt_lb)
test_acc, test_nmi = _evaluate(tt_preds, t_test)
time_test += time.time()
logger.info(
'>>> TEST EPOCH {} ({:.1f}s) lb {:2f} accuracy: {:.2f}% '
'nmi score = {:.4f}'.format(epoch, time_test,
np.mean(tt_lbs),
100. * test_acc, test_nmi))
if epoch == epoches:
with open('results/mnist_scdc.txt', "a") as myfile:
myfile.write("seed: %d train_acc: %f train_nmi: %f "
"test_acc: %f test_nmi: %f" %
(flgs.seed, train_acc, train_nmi,
test_acc, test_nmi))
myfile.write('\n')
myfile.close()
if epoch % save_freq == 0:
images = sess.run(x_gen)
name = os.path.join(result_path, "epoch.{}.png".format(epoch))
utils.save_image_collections(images, name, shape=(10, n_z))