class Test:
def __init__(self, batch_size, data_path, model_path, output_path):
self.epoch = 1
self.batch_size = batch_size
self.model = Model(batch_size)
self.test_dataloader = Dataset(os.path.join(data_path, 'test'))
self.model_path = model_path
self.output_path = output_path
def test(self):
idx_test = 0
block_mask, inverse_block_mask = create_block_mask(shape=[self.batch_size, 96, 96, 3])
with tf.device('/gpu:0'):
with tf.Graph().as_default():
#input
random = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, 96, 96, 3])
y = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, 96, 96, 3])
mask = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, 96, 96, 3])
inverse_mask = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, 96, 96, 3])
x = y*mask + random*inverse_mask
#generator
G_output_sample = self.model.generator(x)
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.6
with tf.Session(config=config) as sess:
sess.run(tf.global_variables_initializer())
restorer = tf.train.Saver()
try:
restorer.restore(sess, tf.train.latest_checkpoint(self.model_path))
print('Load model Success')
except:
print('No model to restore ...')
raise
print('Start testing ...')
count = 0
while True:
epoch, idx_test, y_batch = self.test_dataloader.load_batch(self.batch_size, idx_test, size=[96, 96])
random_noise = np.random.normal(size=y_batch.shape)
G_output_out, x_batch = sess.run([G_output_sample, x], feed_dict={random:random_noise, y:y_batch, mask:block_mask, inverse_mask:inverse_block_mask})
#visualize
G_output_out = G_output_out*inverse_block_mask + y_batch*block_mask
G_output_out = np.squeeze(G_output_out)
print('Saving image {0}' .format(str(count)+'-noiseX-test'))
plot(x_batch, name=str(count)+'-noiseX-test' ,output_path=self.output_path)
print('Saving image {0}' .format(str(count)+'-realX-test'))
plot(y_batch, name=str(count)+'-realX-test' ,output_path=self.output_path)
print('Saving image {0}' .format(str(count)+'-fakeG-test'))
plot(G_output_out, name=str(count)+'-fakeG-test' ,output_path=self.output_path)
count += 1
if epoch == self.epoch:
print('Finish ...')
break
def test():
epoch = 1
batch_size = args.batch_size
model_path = args.model_path
data_path = args.data_path
idx = 0
confusion_matrix_total = np.zeros((10, 10))
data_mode = args.data_mode
data = Dataset(data_mode, data_path, is_shuffle=False)
with tf.Graph().as_default():
x, y, y_onehot, logits = Model(batch_size).build_model('test',
keep_rate=1)
prediction = tf.argmax(logits, axis=1)
accuracy, accuracy_update = tf.metrics.accuracy(labels=y,
predictions=prediction)
batch_confusion = tf.confusion_matrix(labels=y,
predictions=prediction,
num_classes=10)
with tf.Session() as sess:
sess.run([
tf.global_variables_initializer(),
tf.local_variables_initializer()
])
restorer = tf.train.Saver()
try:
restorer.restore(sess, tf.train.latest_checkpoint(model_path))
except:
print('No model in ' + model_path + ' to restore')
raise
while True:
epoch_now, idx, imgs, labels = data.load_batch(batch_size, idx)
_, batch_confusion_ = sess.run(
[accuracy_update, batch_confusion],
feed_dict={
x: imgs,
y: labels
})
confusion_matrix_total += batch_confusion_
if epoch_now == epoch:
np.save('./confusion.npy', confusion_matrix_total)
break
accuracy_ = sess.run(accuracy)
return accuracy_
def train():
idx = 0
batch_size = args.batch_size
learning_rate_start = 0.001
epoch = args.epoch
model_path = args.model_path
data_path = args.data_path
is_restore = args.is_restore
data = Dataset('train', data_path)
with tf.Graph().as_default():
x, y, y_onehot, logits = Model(batch_size).build_model('train', keep_rate=0.5)
cross_entropy = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(labels=y_onehot, logits=logits, name='softmax_loss'))
reg_loss = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)
loss = tf.reduce_sum(cross_entropy) + tf.reduce_sum(reg_loss)
global_step = tf.Variable(0, trainable=False)
learning_rate = tf.train.exponential_decay(learning_rate_start, global_step, 10000, 0.1)
optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
train_op = optimizer.minimize(loss, global_step=global_step)
tf.summary.scalar('loss', loss)
summary = tf.summary.merge_all()
with tf.Session() as sess:
summary_writer = tf.summary.FileWriter('./graph', sess.graph)
sess.run(tf.global_variables_initializer())
saver = tf.train.Saver(max_to_keep=5)
if is_restore:
try:
saver.restore(sess, tf.train.latest_checkpoint(model_path))
except:
print('No model in ' + model_path + ' to restore')
raise
print('Start training...')
while True:
epoch_now, idx, imgs, labels = data.load_batch(batch_size, idx)
loss_, _, step, summary_ = sess.run([loss, train_op, global_step, summary], feed_dict={x: imgs, y:labels})
summary_writer.add_summary(summary_, global_step=step)
if step%20 == 0:
print('Epoch: {0}, Step: {1}, Loss: {2}' .format(epoch_now, step, loss_))
if step%200 == 0:
saver.save(sess, os.path.join(model_path, 'model.ckpt'), global_step=step)
if epoch_now == epoch:
print('Finish training ...')
break
class Train:
def __init__(self,
epoch,
batch_size,
data_path,
model_path,
output_path,
graph_path,
restore=False):
self.batch_size = batch_size
self.model = Model(batch_size)
self.train_dataloader = Dataset(os.path.join(data_path, 'train'))
self.train_test_dataloader = Dataset(os.path.join(data_path, 'train'))
self.test_dataloader = Dataset(os.path.join(data_path, 'test'))
self.epoch = epoch
self.model_path = model_path
self.output_path = output_path
self.graph_path = graph_path
self.restore = restore
def train(self):
idx_train = 0
idx_train_test = 0
idx_test = 0
block_mask, inverse_block_mask = create_block_mask(
shape=[self.batch_size, 96, 96, 3])
with tf.device('/gpu:0'):
with tf.Graph().as_default():
#input
random = tf.placeholder(dtype=tf.float32,
shape=[self.batch_size, 96, 96, 3])
y = tf.placeholder(dtype=tf.float32,
shape=[self.batch_size, 96, 96, 3])
mask = tf.placeholder(dtype=tf.float32,
shape=[self.batch_size, 96, 96, 3])
inverse_mask = tf.placeholder(
dtype=tf.float32, shape=[self.batch_size, 96, 96, 3])
x = y * mask + random * inverse_mask
#generator
G_output = self.model.generator(x)
G_output_sample = self.model.generator(x, reuse=True)
G_output = G_output * inverse_mask + y * mask
#loss
texture_loss = tf.nn.l2_loss(y * mask - G_output * mask)
crop_loss = tf.nn.l2_loss(y * inverse_mask -
G_output * inverse_mask)
X_loss = texture_loss + crop_loss
global_step = tf.Variable(0, trainable=False)
#optimizer
X_op = tf.train.AdamOptimizer(learning_rate=1e-4)
X_train_op = X_op.minimize(X_loss, global_step=global_step)
#tensorboard
tf.summary.scalar('X_loss', X_loss)
summary_op = tf.summary.merge_all()
#GPU ratio
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.6
with tf.Session(config=config) as sess:
sess.run(tf.global_variables_initializer())
saver = tf.train.Saver(max_to_keep=5)
if self.restore:
try:
saver.restore(
sess,
tf.train.latest_checkpoint(self.model_path))
print('Load model Success')
except:
print('No model to restore ...')
raise
summary_writer = tf.summary.FileWriter(
self.graph_path, sess.graph)
print('Start training ...')
while True:
epoch, idx_train, y_batch = self.train_dataloader.load_batch(
self.batch_size, idx_train, size=[96, 96])
random_noise = np.random.normal(size=y_batch.shape)
_, loss_X, step, summary = sess.run(
[X_train_op, X_loss, global_step, summary_op],
feed_dict={
random: random_noise,
y: y_batch,
mask: block_mask,
inverse_mask: inverse_block_mask
})
summary_writer.add_summary(summary, global_step=step)
if step % 10 == 0:
print('Epoch: {0} Step: {1} Loss_X: {2}'.format(
epoch, step, loss_X))
#sample
if step % 100 == 0:
#sample form training data
_, idx_test, y_batch = self.test_dataloader.load_batch(
self.batch_size, idx_test, size=[96, 96])
random_noise = np.random.normal(size=y_batch.shape)
G_output_out, x_batch = sess.run(
[G_output_sample, x],
feed_dict={
random: random_noise,
y: y_batch,
mask: block_mask,
inverse_mask: inverse_block_mask
})
#visualize
G_output_out = G_output_out * inverse_block_mask + y_batch * block_mask
G_output_out = np.squeeze(G_output_out)
plot(x_batch,
name=str(step) + '-noiseX',
output_path=self.output_path)
plot(y_batch,
name=str(step) + '-realX',
output_path=self.output_path)
plot(G_output_out,
name=str(step) + '-fakeG',
output_path=self.output_path)
#sample form testing data
_, idx_train_test, y_batch = self.train_test_dataloader.load_batch(
self.batch_size, idx_train_test, size=[96, 96])
random_noise = np.random.normal(size=y_batch.shape)
G_output_out, x_batch = sess.run(
[G_output_sample, x],
feed_dict={
random: random_noise,
y: y_batch,
mask: block_mask,
inverse_mask: inverse_block_mask
})
#visualize
G_output_out = G_output_out * inverse_block_mask + y_batch * block_mask
G_output_out = np.squeeze(G_output_out)
plot(x_batch,
name=str(step) + '-noiseX-train',
output_path=self.output_path)
plot(y_batch,
name=str(step) + '-realX-train',
output_path=self.output_path)
plot(G_output_out,
name=str(step) + '-fakeG-train',
output_path=self.output_path)
#save model
if step % 1000 == 0:
print('Saving model...')
saver.save(sess,
self.model_path + 'model',
global_step=step)
#finish training
if epoch == self.epoch:
print('Finish training...')
break
class Train:
def __init__(self,
epoch,
batch_size,
data_path,
model_path,
output_path,
graph_path,
restore=False):
self.batch_size = batch_size
self.GAMMA = 0.001
self.Pre_train_content_step = 20000
self.model = Model(batch_size)
self.train_dataloader = Dataset(os.path.join(data_path, 'train'))
self.train_test_dataloader = Dataset(os.path.join(data_path, 'train'))
self.test_dataloader = Dataset(os.path.join(data_path, 'test'))
self.epoch = epoch
self.model_path = model_path
self.output_path = output_path
self.graph_path = graph_path
self.restore = restore
def train(self):
idx_train = 0
idx_train_test = 0
idx_test = 0
with tf.device('/gpu:0'):
with tf.Graph().as_default():
#input
random = tf.placeholder(dtype=tf.float32,
shape=[self.batch_size, 96, 96, 3])
y = tf.placeholder(dtype=tf.float32,
shape=[self.batch_size, 96, 96, 3])
mask = tf.placeholder(dtype=tf.float32,
shape=[self.batch_size, 96, 96, 3])
inverse_mask = tf.placeholder(
dtype=tf.float32, shape=[self.batch_size, 96, 96, 3])
x = y * mask + random * inverse_mask
p_order = tf.random_uniform(shape=[],
minval=0.,
maxval=1.,
dtype=tf.float32)
pred = tf.less(p_order, 0.2)
y_ = tf.cond(pred, lambda: tf.random_shuffle(y), lambda: y)
real_label = tf.cond(pred, lambda: tf.zeros(
(self.batch_size, 1)), lambda: tf.ones(
(self.batch_size, 1)))
#generator
G_output = self.model.generator(x)
G_output_sample = self.model.generator(x, reuse=True)
#discriminator
G_output_mix = y * mask + G_output * inverse_mask
D_real = self.model.discriminator(y_)
D_fake = self.model.discriminator(G_output_mix, reuse=True)
#variable_list
t_vars = tf.trainable_variables()
var_G = [var for var in t_vars if 'G' in var.name]
var_D = [var for var in t_vars if 'D' in var.name]
#loss
#real_label = tf.ones((self.batch_size, 1))
fake_label = tf.zeros((self.batch_size, 1))
D_real_loss = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(labels=real_label,
logits=D_real))
D_fake_loss = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(labels=fake_label,
logits=D_fake))
G_fake_loss = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(labels=real_label,
logits=D_fake))
C_loss = tf.nn.l2_loss(G_output - y)
D_loss = (D_real_loss + D_fake_loss) / 2
G_loss = G_fake_loss + self.GAMMA * C_loss
global_step = tf.Variable(0, trainable=False)
lr = tf.train.exponential_decay(1e-4,
global_step,
100000,
0.1,
staircase=True,
name=None)
#optimizer
G_op = tf.train.AdamOptimizer(learning_rate=lr, beta1=0.5)
D_op = tf.train.AdamOptimizer(learning_rate=lr, beta1=0.5)
D_train_op = D_op.minimize(D_loss,
global_step=global_step,
var_list=var_D)
G_train_op = G_op.minimize(G_loss,
global_step=global_step,
var_list=var_G)
#tensorboard
C_loss_op = tf.summary.scalar('Content_loss', C_loss)
tf.summary.scalar('Generator_loss', G_loss)
tf.summary.scalar('Discriminator_loss', D_loss)
summary_op = tf.summary.merge_all()
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.7
with tf.Session(config=config) as sess:
sess.run(tf.global_variables_initializer())
saver = tf.train.Saver(max_to_keep=5)
if self.restore:
try:
saver.restore(
sess,
tf.train.latest_checkpoint(self.model_path))
except:
print('No model to restore ...')
raise
summary_writer = tf.summary.FileWriter(
self.graph_path, sess.graph)
print('Start training ...')
step = 0
while True:
epoch, idx_train, y_batch = self.train_dataloader.load_batch(
self.batch_size, idx_train, size=[96, 96])
block_mask, inverse_block_mask = create_noise_mask(
y_batch.shape)
random_noise = np.random.normal(size=y_batch.shape)
#Discriminator
if step < self.Pre_train_content_step:
_, loss_C, step, summary = sess.run(
[G_train_op, C_loss, global_step, C_loss_op],
feed_dict={
random: random_noise,
y: y_batch,
mask: block_mask,
inverse_mask: inverse_block_mask
})
summary_writer.add_summary(summary,
global_step=step)
if step % 10 == 0:
print(
'Epoch: {0} Step: {1} Loss_C: {2}'.format(
epoch, step, loss_C))
else:
_, loss_D, loss_G, step, summary = sess.run(
[
D_train_op, D_loss, G_loss, global_step,
summary_op
],
feed_dict={
random: random_noise,
y: y_batch,
mask: block_mask,
inverse_mask: inverse_block_mask
})
summary_writer.add_summary(summary,
global_step=step)
_, loss_D, loss_G, step, summary = sess.run(
[
G_train_op, D_loss, G_loss, global_step,
summary_op
],
feed_dict={
random: random_noise,
y: y_batch,
mask: block_mask,
inverse_mask: inverse_block_mask
})
summary_writer.add_summary(summary,
global_step=step)
if step % 10 == 0:
print(
'Epoch: {0} Step: {1} Loss_D: {2} Loss_G: {3}'
.format(epoch, step, loss_D, loss_G))
#sample
if step % 300 == 0:
#sample training data
_, idx_test, y_batch = self.test_dataloader.load_batch(
self.batch_size, idx_test, size=[96, 96])
block_mask, inverse_block_mask = create_noise_mask(
y_batch.shape)
random_noise = np.random.normal(size=y_batch.shape)
G_output_out, x_batch = sess.run(
[G_output_sample, x],
feed_dict={
random: random_noise,
y: y_batch,
mask: block_mask,
inverse_mask: inverse_block_mask
})
#visualize
G_output_out = G_output_out * inverse_block_mask + y_batch * block_mask
plot(x_batch,
name=str(step) + '-noiseX',
output_path=self.output_path)
plot(y_batch,
name=str(step) + '-realX',
output_path=self.output_path)
plot(G_output_out,
name=str(step) + '-fakeG',
output_path=self.output_path)
#sample testing data
_, idx_train_test, y_batch = self.train_test_dataloader.load_batch(
self.batch_size, idx_train_test, size=[96, 96])
block_mask, inverse_block_mask = create_noise_mask(
y_batch.shape)
random_noise = np.random.normal(size=y_batch.shape)
G_output_out, x_batch = sess.run(
[G_output_sample, x],
feed_dict={
random: random_noise,
y: y_batch,
mask: block_mask,
inverse_mask: inverse_block_mask
})
#visualize
G_output_out = G_output_out * inverse_block_mask + y_batch * block_mask
plot(x_batch,
name=str(step) + '-noiseX-train',
output_path=self.output_path)
plot(y_batch,
name=str(step) + '-realX-train',
output_path=self.output_path)
plot(G_output_out,
name=str(step) + '-fakeG-train',
output_path=self.output_path)
#save model
if step % 1000 == 0:
print('Saving model...')
saver.save(sess,
self.model_path + 'model',
global_step=step)
print('Saving Success')
#finish training
if epoch == self.epoch:
print('Finish training...')
break