with tf.name_scope("KL_divergence"):

loss = -log_sigma + .5 * (-1 + tf.exp(2. * log_sigma) + tf.square(mu))

loss = tf.reduce_mean(loss)

def __init__(self,

model,

dataset=None,

exp_name="model",

ckt_logs_dir="ckt_logs",

):

"""

:type model: RegularizedGAN

"""

self.model = model

self.dataset = dataset

self.exp_name = exp_name

self.log_dir = ckt_logs_dir

self.checkpoint_dir = ckt_logs_dir

self.batch_size = cfg.TRAIN.BATCH_SIZE

self.max_epoch = cfg.TRAIN.MAX_EPOCH

self.snapshot_interval = cfg.TRAIN.SNAPSHOT_INTERVAL

self.model_path = cfg.TRAIN.PRETRAINED_MODEL

self.log_vars = []

def build_placeholder(self):

'''Helper function for init_opt'''

self.images = tf.placeholder(

tf.float32, [self.batch_size] + self.dataset.image_shape,

name='real_images')

self.wrong_images = tf.placeholder(

tf.float32, [self.batch_size] + self.dataset.image_shape,

name='wrong_images'

)

self.conditions = tf.placeholder(

tf.float32, [self.batch_size] + self.dataset.image_shape,

name='conditions'

)

self.embeddings = tf.placeholder(

tf.float32, [self.batch_size] + self.dataset.embedding_shape,

name='conditional_embeddings'

)

self.generator_lr = tf.placeholder(

tf.float32, [],

name='generator_learning_rate'

)

self.discriminator_lr = tf.placeholder(

tf.float32, [],

name='discriminator_learning_rate'

)

# def sample_encoded_context(self, conditions):

# '''Helper function for init_opt'''

# c_mean_logsigma = self.model.generate_condition(conditions)

# mean = c_mean_logsigma[0]

# if cfg.TRAIN.COND_AUGMENTATION:

# # epsilon = tf.random_normal(tf.shape(mean))

# epsilon = tf.truncated_normal(tf.shape(mean))

# stddev = tf.exp(c_mean_logsigma[1])

# c = mean + stddev * epsilon

# kl_loss = KL_loss(c_mean_logsigma[0], c_mean_logsigma[1])

# else:

# c = mean

# kl_loss = 0

# return c, cfg.TRAIN.COEFF.KL * kl_loss

def init_opt(self):

self.build_placeholder()

with tf.variable_scope("g_net"):

with pt.defaults_scope(phase=pt.Phase.train):

# ####get output from G network################################

# c, kl_loss = self.sample_encoded_context(self.conditions)

#c_text = self.model.generate_text_condition(self.embeddings)

#c, c_vector = self.model.generate_condition(self.conditions)

c_text = self.embeddings

c_mask = self.conditions

z = tf.random_normal([self.batch_size, cfg.Z_DIM])

self.log_vars.append(("hist_c", c_text))

self.log_vars.append(("hist_z", z))

#G_input = tf.concat(1, [c, c_text])

#G_input = tf.concat(1, [G_input, z])

# print('#######################')

# print('c_text, c_mask shape is:', c_text, c_mask)

# print('#########################')

G_input = self.model.generate_condition(tf.concat(1, [c_text, z]), c_mask)

# update by zhaoyang

self.debug_g_input = G_input

###################

fake_images = self.model.get_generator(G_input)

# update by kristy

self.debug_fake_images = fake_images

# ####get discriminator_loss and generator_loss ###################

discriminator_loss, generator_loss =\

self.compute_losses(self.images,

self.wrong_images,

fake_images,

self.embeddings,

self.conditions)

#generator_loss += kl_loss

#self.log_vars.append(("g_loss_kl_loss", kl_loss))

self.log_vars.append(("g_loss", generator_loss))

self.log_vars.append(("d_loss", discriminator_loss))

# #######Total loss for build optimizers###########################

self.prepare_trainer(generator_loss, discriminator_loss)

# #######define self.g_sum, self.d_sum,....########################

self.define_summaries()

with pt.defaults_scope(phase=pt.Phase.test):

#with tf.variable_scope("g_net", reuse=True):

self.sampler()

self.visualization(cfg.TRAIN.NUM_COPY)

print("success")

def sampler(self):

# c, _ = self.sample_encoded_context(self.conditions)

c_text = self.embeddings

c_mask = self.conditions

#if cfg.TRAIN.FLAG:

# z = tf.zeros([self.batch_size, cfg.Z_DIM]) # Expect similar BGs

#else:

# z = tf.random_normal([self.batch_size, cfg.Z_DIM])

z = tf.random_normal([self.batch_size, cfg.Z_DIM])

G_input = self.model.generate_condition(tf.concat(1, [c_text, z]), c_mask)

self.fake_images = self.model.get_generator(G_input)

def compute_losses(self, images, wrong_images, fake_images, embeddings, conditions):

real_logit = self.model.get_discriminator(images, embeddings, conditions)

wrong_logit = self.model.get_discriminator(wrong_images, embeddings, conditions)

fake_logit = self.model.get_discriminator(fake_images, embeddings, conditions)

real_d_loss =\

tf.nn.sigmoid_cross_entropy_with_logits(real_logit,

tf.ones_like(real_logit))

real_d_loss = tf.reduce_mean(real_d_loss)

wrong_d_loss =\

tf.nn.sigmoid_cross_entropy_with_logits(wrong_logit,

tf.zeros_like(wrong_logit))

wrong_d_loss = tf.reduce_mean(wrong_d_loss)

fake_d_loss =\

tf.nn.sigmoid_cross_entropy_with_logits(fake_logit,

tf.zeros_like(fake_logit))

fake_d_loss = tf.reduce_mean(fake_d_loss)

if cfg.TRAIN.B_WRONG:

discriminator_loss =\

real_d_loss + (wrong_d_loss + fake_d_loss) / 2.

self.log_vars.append(("d_loss_wrong", wrong_d_loss))

else:

discriminator_loss = real_d_loss + fake_d_loss

self.log_vars.append(("d_loss_real", real_d_loss))

self.log_vars.append(("d_loss_fake", fake_d_loss))

generator_loss = \

tf.nn.sigmoid_cross_entropy_with_logits(fake_logit,

tf.ones_like(fake_logit))

generator_loss = tf.reduce_mean(generator_loss)

return discriminator_loss, generator_loss

def prepare_trainer(self, generator_loss, discriminator_loss):

'''Helper function for init_opt'''

all_vars = tf.trainable_variables()

g_vars = [var for var in all_vars if

var.name.startswith('g_')]

d_vars = [var for var in all_vars if

var.name.startswith('d_')]

generator_opt = tf.train.AdamOptimizer(self.generator_lr,

beta1=0.5)

self.generator_trainer =\

pt.apply_optimizer(generator_opt,

losses=[generator_loss],

var_list=g_vars)

discriminator_opt = tf.train.AdamOptimizer(self.discriminator_lr,

beta1=0.5)

self.discriminator_trainer =\

pt.apply_optimizer(discriminator_opt,

losses=[discriminator_loss],

var_list=d_vars)

self.log_vars.append(("g_learning_rate", self.generator_lr))

self.log_vars.append(("d_learning_rate", self.discriminator_lr))

def define_summaries(self):

'''Helper function for init_opt'''

all_sum = {'g': [], 'd': [], 'hist': []}

for k, v in self.log_vars:

if k.startswith('g'):

all_sum['g'].append(tf.scalar_summary(k, v))

elif k.startswith('d'):

all_sum['d'].append(tf.scalar_summary(k, v))

elif k.startswith('hist'):

all_sum['hist'].append(tf.histogram_summary(k, v))

self.g_sum = tf.merge_summary(all_sum['g'])

self.d_sum = tf.merge_summary(all_sum['d'])

self.hist_sum = tf.merge_summary(all_sum['hist'])

def visualize_one_superimage(self, img_var, images, rows, filename):

stacked_img = []

for row in range(rows):

img = images[row * rows, :, :, :]

row_img = [img] # real image

for col in range(rows):

row_img.append(img_var[row * rows + col, :, :, :])

# each rows is 1realimage +10_fakeimage

stacked_img.append(tf.concat(1, row_img))

imgs = tf.expand_dims(tf.concat(0, stacked_img), 0)

current_img_summary = tf.image_summary(filename, imgs)

return current_img_summary, imgs

def visualization(self, n):

fake_sum_train, superimage_train = \

self.visualize_one_superimage(self.debug_fake_images[:n * n],

self.images[:n * n],

n, "train")

fake_sum_test, superimage_test = \

self.visualize_one_superimage(self.debug_fake_images[n * n:2 * n * n],

self.images[n * n:2 * n * n],

n, "test")

self.superimages = tf.concat(0, [superimage_train, superimage_test])

self.image_summary = tf.merge_summary([fake_sum_train, fake_sum_test])

def preprocess(self, x, n):

# make sure every row with n column have the same conditions

for i in range(n):

for j in range(1, n):

x[i * n + j] = x[i * n]

return x

def epoch_sum_images(self, sess, n, counter):

images_train, _, conditions_train, embeddings_train, captions_train, _ =\

self.dataset.train.next_batch(n * n, cfg.TRAIN.NUM_EMBEDDING)

images_train = self.preprocess(images_train, n)

conditions_train = self.preprocess(conditions_train, n)

embeddings_train = self.preprocess(embeddings_train, n)

images_test, _, conditions_test, embeddings_test, captions_test, _ = \

self.dataset.test.next_batch(n * n, 1)

images_test = self.preprocess(images_test, n)

conditions_test = self.preprocess(conditions_test, n)

embeddings_test = self.preprocess(embeddings_test, n)

images = np.concatenate([images_train, images_test], axis=0)

conditions =\

np.concatenate([conditions_train, conditions_test], axis=0)

embeddings =\

np.concatenate([embeddings_train, embeddings_test], axis=0)

if self.batch_size > 2 * n * n:

images_pad, _, conditions_pad, embeddings_pad, _, _ =\

self.dataset.test.next_batch(self.batch_size - 2 * n * n, 1)

#print ('embedding and embedding pad shape is:', embeddings.shape, embeddings_pad.shape)

images = np.concatenate([images, images_pad], axis=0)

conditions = np.concatenate([conditions, conditions_pad], axis=0)

embeddings = np.concatenate([embeddings, embeddings_pad], axis=0)

feed_dict = {self.images: images,

self.conditions: conditions,

self.embeddings: embeddings}

fake, gen_samples, img_summary =\

sess.run([self.debug_fake_images[0], self.superimages, self.image_summary], feed_dict)

#print('fake image is:', fake)

#print('gen_sampls[0]', gen_samples[0])

# save images generated for train and test captions

#scipy.misc.imsave('%s/fake.jpg' % (self.log_dir), fake)

scipy.misc.imsave('%s/train_%d.jpg' % (self.log_dir, counter), gen_samples[0])

scipy.misc.imsave('%s/test_%d.jpg' % (self.log_dir, counter), gen_samples[1])

########## added by kristy -- visualize the seg condition#######################

# print ('conditions shape is', conditions.shape)

# scipy.misc.imsave('%s/condition_test.jpg' % (self.log_dir), conditions[0])

# pfi_train = open(self.log_dir + "/train.txt", "w")

pfi_test = open(self.log_dir + "/test_%d.txt" %counter, "w")

for row in range(n):

# pfi_train.write('\n***row %d***\n' % row)

# pfi_train.write(captions_train[row * n])

pfi_test.write('\n***row %d***\n' % row)

pfi_test.write(captions_test[row * n])

# pfi_train.close()

pfi_test.close()

return img_summary

def build_model(self, sess):

self.init_opt()

sess.run(tf.initialize_all_variables())

if len(self.model_path) > 0:

print("Reading model parameters from %s" % self.model_path)

restore_vars = tf.all_variables()

# all_vars = tf.all_variables()

# restore_vars = [var for var in all_vars if

# var.name.startswith('g_') or

# var.name.startswith('d_')]

saver = tf.train.Saver(restore_vars)

saver.restore(sess, self.model_path)

istart = self.model_path.rfind('_') + 1

iend = self.model_path.rfind('.')

counter = self.model_path[istart:iend]

counter = int(counter)

else:

print("Created model with fresh parameters.")

counter = 0

return counter

def train(self):

config = tf.ConfigProto(allow_soft_placement=True)

with tf.Session(config=config) as sess:

#with tf.device("/gpu:%d" % cfg.GPU_ID):

counter = self.build_model(sess)

saver = tf.train.Saver(tf.all_variables(),

keep_checkpoint_every_n_hours=2)

# summary_op = tf.merge_all_summaries()

summary_writer = tf.train.SummaryWriter(self.log_dir,

sess.graph)

keys = ["d_loss", "g_loss"]

log_vars = []

log_keys = []

for k, v in self.log_vars:

if k in keys:

log_vars.append(v)

log_keys.append(k)

# print(k, v)

generator_lr = cfg.TRAIN.GENERATOR_LR

discriminator_lr = cfg.TRAIN.DISCRIMINATOR_LR

num_embedding = cfg.TRAIN.NUM_EMBEDDING

lr_decay_step = cfg.TRAIN.LR_DECAY_EPOCH

number_example = self.dataset.train._num_examples

updates_per_epoch = int(number_example / self.batch_size)

epoch_start = int(counter / updates_per_epoch)

for epoch in range(epoch_start, self.max_epoch):

widgets = ["epoch #%d|" % epoch,

Percentage(), Bar(), ETA()]

pbar = ProgressBar(maxval=updates_per_epoch,

widgets=widgets)

pbar.start()

if epoch % lr_decay_step == 0 and epoch != 0:

generator_lr *= 0.5

discriminator_lr *= 0.5

all_log_vals = []

for i in range(updates_per_epoch):

pbar.update(i)

# training d

images, wrong_images, conditions, embeddings, _, _ =\

self.dataset.train.next_batch(self.batch_size, num_embedding)

# print('##########saving images and conditions ...######')

# scipy.misc.imsave('image_%d.jpg' %counter, images[0])

# scipy.misc.imsave('wrong_images_%d.jpg' %counter, wrong_images[0])

# scipy.misc.imsave('condition_%d.jpg' %counter, conditions[0])

# print('##############complete!########################')

feed_dict = {self.images: images,

self.wrong_images: wrong_images,

self.conditions: conditions,

self.embeddings: embeddings,

self.generator_lr: generator_lr,

self.discriminator_lr: discriminator_lr

}

# train d

feed_out = [self.discriminator_trainer,

self.d_sum,

self.hist_sum,

log_vars]

_, d_sum, hist_sum, log_vals = sess.run(feed_out,

feed_dict)

# feed_out = [self.discriminator_trainer, log_vars]

# _, log_vals = sess.run(feed_out,feed_dict)

summary_writer.add_summary(d_sum, counter)

summary_writer.add_summary(hist_sum, counter)

all_log_vals.append(log_vals)

# train g

feed_out = [self.generator_trainer,

self.g_sum]

_, g_sum = sess.run(feed_out,

feed_dict)

# feed_out = [self.generator_trainer]

# _, =sess.run(feed_out,

# feed_dict)

summary_writer.add_summary(g_sum, counter)

# # update by zhaoyang###############################

# # this part of code is very very very important

# # feed_out = [self.discriminator_trainer, log_vars, self.debug_g_input]

# # _, log_vals, debug_output = sess.run(feed_out,feed_dict)

#feed_out = [self.debug_g_input, self.fake_images]

#debug_output, debug_img = sess.run(feed_out,feed_dict)

#fake_pic = np.empty(debug_img[0].shape)

#fake_pic = debug_img[0]

#import scipy.misc

#print('debug info:########')

# #print('G_input is :', debug_output)

# #print('fake_image is:', debug_img)

#img_name = 'fake_img_%d.jpg' %(counter)

#scipy.misc.imsave(img_name,fake_pic)

#print('save img complete!!!')

#print('########')

# ###################

# save checkpoint

counter += 1

if counter % self.snapshot_interval == 0:

snapshot_path = "%s/%s_%s.ckpt" %\

(self.checkpoint_dir,

self.exp_name,

str(counter))

fn = saver.save(sess, snapshot_path)

print("Model saved in file: %s" % fn)

img_sum = self.epoch_sum_images(sess, cfg.TRAIN.NUM_COPY, counter)

#summary_writer.add_summary(img_sum, counter)

avg_log_vals = np.mean(np.array(all_log_vals), axis=0)

dic_logs = {}

for k, v in zip(log_keys, avg_log_vals):

dic_logs[k] = v

# print(k, v)

log_line = "; ".join("%s: %s" %

(str(k), str(dic_logs[k]))

for k in dic_logs)

print("Epoch %d | " % (epoch) + log_line)

sys.stdout.flush()

if np.any(np.isnan(avg_log_vals)):

raise ValueError("NaN detected!")

def save_super_images(self, images, sample_batchs, filenames,

sentenceID, save_dir, subset):

# batch_size samples for each embedding

numSamples = len(sample_batchs)

for j in range(len(filenames)):

s_tmp = '%s-1real-%dsamples/%s/%s' %\

(save_dir, numSamples, subset, filenames[j])

folder = s_tmp[:s_tmp.rfind('/')]

if not os.path.isdir(folder):

print('Make a new folder: ', folder)

mkdir_p(folder)

superimage = [images[j]]

# cfg.TRAIN.NUM_COPY samples for each text embedding/sentence

for i in range(len(sample_batchs)):

superimage.append(sample_batchs[i][j])

superimage = np.concatenate(superimage, axis=1)

fullpath = '%s_sentence%d.jpg' % (s_tmp, sentenceID)

scipy.misc.imsave(fullpath, superimage)

def eval_one_dataset(self, sess, dataset, save_dir, subset='train'):

count = 0

print('num_examples:', dataset._num_examples)

while count < dataset._num_examples:

start = count % dataset._num_examples

images, conditions_batchs, embeddings_batches,

filenames, _ =\

dataset.next_batch_test(self.batch_size, start, 1)

print('count = ', count, 'start = ', start)

for i in range(len(conditions_batchs)):

samples_batchs = []

# Generate up to 16 images for each sentence,

# with randomness from noise z and conditioning augmentation.

for j in range(np.minimum(16, cfg.TRAIN.NUM_COPY)):

samples = sess.run(self.fake_images,

{self.conditions: conditions_batchs[i]})

samples_batchs.append(samples)

self.save_super_images(images, samples_batchs,

filenames, i, save_dir,

subset)

count += self.batch_size

def evaluate(self):

config = tf.ConfigProto(allow_soft_placement=True)

with tf.Session(config=config) as sess:

#with tf.device("/gpu:%d" % cfg.GPU_ID):

if self.model_path.find('.ckpt') != -1:

self.init_opt(sess)

print("Reading model parameters from %s" % self.model_path)

saver = tf.train.Saver(tf.all_variables())

saver.restore(sess, self.model_path)

# self.eval_one_dataset(sess, self.dataset.train,

# self.log_dir, subset='train')

self.eval_one_dataset(sess, self.dataset.test,

self.log_dir, subset='test')

else: