def run_test_loop(self):
G_sample, G_relectance, G_alpha = self.generator(self.Z, self.BG)
sess = tf.Session()
sess.run(tf.global_variables_initializer())
# load the detection model
self.load_historical_model(sess)
if not os.path.exists(self.result_dir):
os.makedirs(self.result_dir)
for i in range(self.iter_step):
BG_mb = utils.get_batch(self.data,
self.batch_size,
mode='background',
with_data_augmentation=True,
img_size=self.img_size)
Z_mb = utils.get_batch(self.data,
self.batch_size,
mode='cloudimage',
with_data_augmentation=True,
img_size=self.img_size)
sample, reflectance, alpha = sess.run(
[G_sample, G_relectance, G_alpha],
feed_dict={
self.Z: Z_mb,
self.BG: BG_mb
})
for jj in range(self.batch_size):
save_path = os.path.join(
self.result_dir, '{}'.format(str(i).zfill(5)) + '_' +
str(jj) + '_image.jpg')
plt.imsave(save_path, sample[jj, :, :, :])
save_path = os.path.join(
self.result_dir, '{}'.format(str(i).zfill(5)) + '_' +
str(jj) + '_reflectance.png')
plt.imsave(save_path, reflectance[jj, :, :, :])
save_path = os.path.join(
self.result_dir, '{}'.format(str(i).zfill(5)) + '_' +
str(jj) + '_alpha.png')
plt.imsave(save_path, alpha[jj, :, :, :])
def run_train_loop(self):
params = init.TrainingParamInitialization()
self.gan = CloudGAN(params)
self.train_img, self.train_alpha, self.train_reflectance =\
self.data_prepare.preprocess(self.gan.G_sample, self.gan.G_alpha, self.gan.G_relectance)
self.validate_img, self.validate_alpha, self.validate_reflectance = \
self.data_prepare.preprocess(self.gan.G_sample, self.gan.G_alpha,self.gan.G_relectance)
print('build matting net')
with tf.name_scope('train'):
train_op = self.train_op(self.train_img, self.train_alpha,
self.train_reflectance)
with tf.name_scope('validate'):
validate_op = self.validate_op(self.validate_img,
self.validate_alpha,
self.validate_reflectance,
reuse=True)
sess = tf.Session()
sess.run(tf.global_variables_initializer())
self.load_checkpoints(sess)
step = 0
log_step = 50
saver = tf.train.Saver()
merge_op = tf.summary.merge_all()
summary_writer = tf.summary.FileWriter(self.logdir, sess.graph)
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(coord=coord, sess=sess)
if not os.path.exists('rst'):
os.mkdir('rst')
while step < self.iter_step:
feed_dict = None
self.gan.run_train(sess)
X_mb = utils.get_batch(self.gan.data, self.batch_size,
'cloudimage', self.img_size)
Z_mb = utils.get_batch(self.gan.data, self.batch_size,
'cloudimage', self.img_size)
BG_mb = utils.get_batch(self.gan.data, self.batch_size,
'background', self.img_size)
feed_dict = {
self.gan.X: X_mb,
self.gan.Z: Z_mb,
self.gan.BG: BG_mb
}
[_, train_loss, step] = sess.run(train_op, feed_dict=feed_dict)
if np.mod(step, log_step) == 1:
X_mb = utils.get_batch(self.gan.data, self.batch_size,
'cloudimage', self.img_size)
Z_mb = utils.get_batch(self.gan.data, self.batch_size,
'cloudimage', self.img_size)
BG_mb = utils.get_batch(self.gan.data, self.batch_size,
'background', self.img_size)
feed_dict = {
self.gan.X: X_mb,
self.gan.Z: Z_mb,
self.gan.BG: BG_mb
}
merges = [merge_op] + validate_op
summary, image, alpha_image, reflectance_image, validate_loss = sess.run(
merges, feed_dict=feed_dict)
summary_writer.add_summary(summary, step)
print("step:%d,loss:%f,validate_loss:%f" %
(step, train_loss, validate_loss))
misc.imsave('rst/' + str(step) + '_image.png', image)
misc.imsave('rst/' + str(step) + '_alpha.png', alpha_image)
misc.imsave('rst/' + str(step) + '_foreground.png',
reflectance_image)
if np.mod(step, 500) == 1:
saver.save(sess, os.path.join(self.model_path, 'model'), step)
def run_train(self, sess):
X_mb = utils.get_batch(self.data, self.batch_size, 'cloudimage',
self.img_size)
Z_mb = utils.get_batch(self.data, self.batch_size, 'cloudimage',
self.img_size)
BG_mb = utils.get_batch(self.data, self.batch_size, 'background',
self.img_size)
_, D_loss_curr, _ = sess.run([self.D_solver, self.D_loss, self.clip_D],
feed_dict={
self.X: X_mb,
self.Z: Z_mb,
self.BG: BG_mb
})
d_iter = sess.run(self.D_steps)
if not os.path.exists(self.sample_dir):
os.mkdir(self.sample_dir)
mm = 1
if D_loss_curr < 0.1:
mm = mm + 1
else:
mm = 1
for _ in range(mm):
X_mb = utils.get_batch(self.data, self.batch_size, 'cloudimage',
self.img_size)
Z_mb = utils.get_batch(self.data, self.batch_size, 'cloudimage',
self.img_size)
BG_mb = utils.get_batch(self.data, self.batch_size, 'background',
self.img_size)
_, G_loss_curr, D_loss_curr = sess.run(
[self.G_solver, self.G_loss, self.D_loss],
feed_dict={
self.X: X_mb,
self.Z: Z_mb,
self.BG: BG_mb
})
g_iter = sess.run(self.G_steps)
# save generated samples
if g_iter % 50 == 0:
samples, reflectance, alpha = sess.run(
[self.G_sample, self.G_relectance, self.G_alpha],
feed_dict={
self.Z: Z_mb,
self.BG: BG_mb
})
save_path = os.path.join(
self.sample_dir,
'{}_reflectance.png'.format(str(g_iter).zfill(5)))
plt.imsave(save_path,
utils.plot2x2(reflectance),
vmax=1,
vmin=0)
save_path = os.path.join(
self.sample_dir,
'{}_image.png'.format(str(g_iter).zfill(5)))
plt.imsave(save_path, utils.plot2x2(samples), vmax=1, vmin=0)
save_path = os.path.join(
self.sample_dir,
'{}_input.png'.format(str(g_iter).zfill(5)))
plt.imsave(save_path, utils.plot2x2(Z_mb), vmax=1, vmin=0)
save_path = os.path.join(
self.sample_dir,
'{}_alpha.png'.format(str(g_iter).zfill(5)))
plt.imsave(save_path, utils.plot2x2(alpha), vmax=1, vmin=0)
if g_iter % 50 == 0:
print('D_loss = %g, G_loss = %g' % (D_loss_curr, G_loss_curr))
print('g_iter = %d, d_iter = %d, n_g/d = %d' %
(g_iter, d_iter, mm))
print()
def run_train_loop(self):
G_sample, G_relectance, G_alpha = self.generator(self.Z, self.BG)
D_logits_real, D_prob_real = self.discriminator(self.X)
D_logits_fake, D_prob_fake = self.discriminator(G_sample, reuse=True)
G_loss, D_loss = self.calc_loss(D_logits_real, D_logits_fake,
G_relectance)
tvars = tf.trainable_variables()
theta_D = [var for var in tvars if 'Discriminator_scope' in var.name]
theta_G = [var for var in tvars if 'Generator_scope' in var.name]
# to record the iteration number
G_steps = tf.Variable(0, trainable=False, name='G_steps')
D_steps = tf.Variable(0, trainable=False, name='D_steps')
with tf.variable_scope(tf.get_variable_scope(),
reuse=tf.AUTO_REUSE) as scope:
if self.optimizer is 'RMSProp':
D_solver = (tf.train.RMSPropOptimizer(
learning_rate=self.d_learning_rate).minimize(
D_loss, var_list=theta_D, global_step=D_steps))
G_solver = (tf.train.RMSPropOptimizer(
learning_rate=self.g_learning_rate).minimize(
G_loss, var_list=theta_G, global_step=G_steps))
if self.optimizer is 'Adam':
D_solver = (tf.train.AdamOptimizer(
learning_rate=self.d_learning_rate).minimize(
D_loss, var_list=theta_D, global_step=D_steps))
G_solver = (tf.train.AdamOptimizer(
learning_rate=self.g_learning_rate).minimize(
G_loss, var_list=theta_G, global_step=G_steps))
clip_D = [
p.assign(tf.clip_by_value(p, -self.d_clip, self.d_clip))
for p in theta_D
]
sess = tf.Session()
sess.run(tf.global_variables_initializer())
# create a summary writer
merged_op = tf.summary.merge_all()
summary_writer = tf.summary.FileWriter('temp', sess.graph)
# load historical model and create a saver
self.load_historical_model(sess)
saver = tf.train.Saver()
if not os.path.exists(self.sample_dir):
os.makedirs(self.sample_dir)
g_iter = 0
d_iter = 0
mm = 1 # number of G_steps per D_step
while g_iter < self.iter_step:
X_mb = utils.get_batch(self.data, self.batch_size, 'cloudimage',
self.img_size)
Z_mb = utils.get_batch(self.data, self.batch_size, 'cloudimage',
self.img_size)
BG_mb = utils.get_batch(self.data, self.batch_size, 'background',
self.img_size)
_, D_loss_curr, _, summary = sess.run(
[D_solver, D_loss, clip_D, merged_op],
feed_dict={
self.X: X_mb,
self.Z: Z_mb,
self.BG: BG_mb
})
d_iter = sess.run(D_steps)
# write states to summary
summary_writer.add_summary(summary, g_iter)
# To stabilize training, we train multiple steps (mm) on G if D loss is less than a pre-defined
# threshold, say, 0.1. We found this simple mofification on the training config greatly prevents
# from model collapse.
if D_loss_curr < 0.1:
mm = mm + 5
else:
mm = 1
for _ in range(mm):
X_mb = utils.get_batch(self.data, self.batch_size,
'cloudimage', self.img_size)
Z_mb = utils.get_batch(self.data, self.batch_size,
'cloudimage', self.img_size)
BG_mb = utils.get_batch(self.data, self.batch_size,
'background', self.img_size)
_, G_loss_curr, D_loss_curr = sess.run(
[G_solver, G_loss, D_loss],
feed_dict={
self.X: X_mb,
self.Z: Z_mb,
self.BG: BG_mb
})
g_iter = sess.run(G_steps)
# save generated samples
if g_iter % 5 == 0:
samples, reflectance, alpha = sess.run(
[G_sample, G_relectance, G_alpha],
feed_dict={
self.Z: Z_mb,
self.BG: BG_mb
})
samples = np.array(samples * 255, np.uint8)
reflectance = np.array(reflectance * 255, np.uint8)
alpha = np.array(alpha * 255, np.uint8)
BG = np.array(BG_mb * 255, np.uint8)
for jj in range(self.batch_size):
i = g_iter
save_path = os.path.join(
self.sample_dir, '{}'.format(str(i).zfill(5)) +
'_' + str(jj) + '.jpg')
plt.imsave(save_path, samples[jj, :, :, :])
save_path = os.path.join(
self.sample_dir, '{}'.format(str(i).zfill(5)) +
'_' + str(jj) + '_reflectance.png')
plt.imsave(save_path, reflectance[jj, :, :, :])
save_path = os.path.join(
self.sample_dir, '{}'.format(str(i).zfill(5)) +
'_' + str(jj) + '_alpha.png')
plt.imsave(save_path, alpha[jj, :, :, :])
save_path = os.path.join(
self.sample_dir, '{}'.format(str(i).zfill(5)) +
'_' + str(jj) + '_BG.png')
plt.imsave(save_path, BG[jj, :, :, :])
if g_iter % 50 == 0:
print('D_loss = %g, G_loss = %g' %
(D_loss_curr, G_loss_curr))
print('g_iter = %d, d_iter = %d, n_g/d = %d' %
(g_iter, d_iter, mm))
# save model every 500 g_iters
if np.mod(g_iter, 500) == 1 and g_iter > 1:
print('saving model to checkpoint ...')
saver.save(sess,
os.path.join(self.checkpoint_gan, 'G_step'),
global_step=G_steps)
def train():
# Turn on training mode which enables dropout.
total_loss, avrg_loss = 0, 0
start_time = time.time()
ntokens = len(corpus.dictionary)
batch, i = 0, 0
# need a hidden state for tl and one for mos
h_tl = tl_model.init_hidden(args.batch_size)
h_mos = mos_model.init_hidden(args.batch_size)
data_keep = 5 * torch.ones(1, args.batch_size).cuda().long()
while i < train_data.size(0) - 1:
# get seq len from batch
seq_len = seq_lens[batch] - 1
# adapt learning rate
lr2 = optimizer.param_groups[0]['lr']
optimizer.param_groups[0]['lr'] = lr2 * seq_len / args.bptt
tl_model.train()
# get data and binary map
data = get_batch(train_data, i, args, seq_len=seq_len)
binary = get_batch(binary_data, i, args, seq_len=seq_len)
# evaluate mos on data
h_mos = repackage_hidden(h_mos)
mos_data = torch.cat((data_keep, data), 0)[:-1]
mos_data[mos_data >= 10000] = 0 # ugly fix!!!!
log_prob, h_mos = mos_model(mos_data, h_mos)
#print(torch.exp(log_prob))
#print(data, mos_data)
# get probability ranks from mos probability
_, argsort = torch.sort(log_prob, descending=True)
argsort = argsort.view(-1, 10000)
#print(argsort.size())
# Starting each batch, we detach the hidden state from how it was previously produced.
# If we didn't, the model would try backpropagating all the way to start of the dataset.
h_tl = tl_model.init_hidden(args.batch_size)
h_tl = repackage_hidden(h_tl)
optimizer.zero_grad()
#raw_loss = model.train_crossentropy(data, eos_tokens)
raw_loss = tl_model(data, binary, h_tl, argsort)
avrg_loss = avrg_loss + (seq_len +
1) * raw_loss.data / train_data.size(0)
loss = raw_loss
loss.backward()
# `clip_grad_norm` helps prevent the exploding gradient problem in RNNs / LSTMs.
if args.clip: torch.nn.utils.clip_grad_norm_(params, args.clip)
optimizer.step()
total_loss += loss.data
optimizer.param_groups[0]['lr'] = lr2
if batch % args.log_interval == 0:
cur_loss = total_loss.item() / args.log_interval
elapsed = time.time() - start_time
print(
'| epoch {:3d} | {:5d}/{:5d} batches | lr {:05.5f} | ms/batch {:5.2f} | '
'loss {:5.2f} | ppl {:8.2f} | bpc {:8.3f}'.format(
epoch, batch,
len(train_data) // args.bptt,
optimizer.param_groups[0]['lr'],
elapsed * 1000 / args.log_interval, cur_loss,
np.exp(cur_loss), cur_loss / math.log(2)))
total_loss = 0
start_time = time.time()
###
batch += 1
i += seq_len + 1
#break
return avrg_loss #/ train_data.size(0)
def evaluate(data_source, epoch, batch_size=1):
# Turn on evaluation mode which disables dropout.
tl_model.eval()
total_loss, i = 0, 0
h_tl = tl_model.init_hidden(batch_size)
h_mos = mos_model.init_hidden(batch_size)
data_keep = 5 * torch.ones(1, 1).cuda().long()
while i < data_source.size(0) - 1:
seq_len = args.bptt # one to many (data has size 2, need this because targets!)
data = get_batch(data_source, i, args, seq_len=seq_len)
mos_data = torch.cat(
(data_keep, data),
0)[:-1] # no need to include eos token in mos_data
data_keep = data[-1].view(1, 1)
# evaluate mos for probability ranks
h_mos = repackage_hidden(h_mos)
log_prob, h_mos = mos_model(mos_data, h_mos)
#print(torch.exp(log_prob), data.view(-1), mos_data.view(-1))
# cut log_probs to actual sequence length
#log_prob = log_prob[:-1]
#print(data, mos_data)
# get probability ranks from mos probability
#print(torch.exp(log_prob), data[1:])
_, argsort = torch.sort(log_prob, descending=True)
argsort = argsort.view(-1, 10000) #
# evaluate tl model
h_tl = repackage_hidden(h_tl)
loss, h_tl, entropy = tl_model.evaluate(data, h_tl, argsort,
eos_tokens)
total_loss = total_loss + loss * min(seq_len, data_source.size(0))
i = i + seq_len + 1
total_loss = total_loss / data_source.size(0)
'''
if args.dump_hiddens:
loss, entropy, hiddens = tl_model.evaluate(data_source, eos_tokens, args.dump_hiddens)
dump_hiddens(hiddens, 'hiddens_' + str(epoch))
else:
loss, entropy = tl_model.evaluate(data_source, eos_tokens)
#loss = loss.item()
if args.dump_words:
W = tl_model.rnn.module.weight_ih_l0.detach()
dump_words(torch.nn.functional.linear(tl_model.encoder.weight.detach(), W).detach().cpu().numpy(), 'words_xW_' + str(epoch))
dump_words(tl_model.encoder.weight.detach().cpu().numpy(), 'words_' + str(epoch))
if not args.dump_entropy is None:
dump(entropy, args.dump_entropy + str(epoch))
'''
return total_loss