def train(train_dataset, args, lsgan=True):
for iteration in range(args.start_iter, args.max_iterations):
start = time.time()
with tf.GradientTape() as genA2B_tape, tf.GradientTape() as discB_tape:
try:
train_full = next(train_dataset)
trainA_full = train_full[:, :, :1024, :]
trainB_full = train_full[:, :, 1024:, :]
except tf.errors.OutOfRangeError:
print("Error, run out of data")
break
split_list_A = pyramid.split(trainA_full, args.level)
trainA = split_list_A[-1]
genA2B_output = genA2B(trainA, training=True)
genA2B_output_full = merge_image(genA2B_output, split_list_A)
discB_real_output = discB(trainB_full, training=True)
discB_fake_output = discB(genA2B_output_full, training=True)
# Use history buffer of 50 for disc loss
discB_loss = discriminator_loss(discB_real_output,
discB_fake_output,
lsgan=lsgan)
generatorA2B_loss = generator_loss(discB_fake_output, lsgan=lsgan)
mse_loss = tf.reduce_mean(
tf.square(trainB_full - genA2B_output_full))
genA2B_loss = args.gan_lambda * generatorA2B_loss + args.mse_lambda * mse_loss
genA2B_gradients = genA2B_tape.gradient(genA2B_loss,
genA2B.trainable_variables)
discB_gradients = discB_tape.gradient(discB_loss,
discB.trainable_variables)
genA2B_optimizer.apply_gradients(
zip(genA2B_gradients, genA2B.trainable_variables))
discB_optimizer.apply_gradients(
zip(discB_gradients, discB.trainable_variables))
if iteration % args.log_interval == 0:
# print('Training Iteration: {}th, LOSSES: D_A: {}, D_B: {}, G_A: {}, G_B: {}, cycle: {}'.format(
# iteration, discA_loss, discB_loss, generatorA2B_loss, generatorB2A_loss, cycleA2B_loss))
print('Training Iteration: {}th, LOSSES: mse: {:.4f}'.format(
iteration, mse_loss))
if iteration % args.save_interval == 0:
generate_images(trainA_full, trainB_full, genA2B_output_full,
genA2B_output_full, args.save_dir_train, iteration)
# print('Time taken for iteration {} is {} sec'.format(iteration + 1, time.time() - start))
if not os.path.exists(args.model_save_dir):
os.mkdir(args.model_save_dir)
discB.save_weights(args.model_save_dir + '/discB_' +
str(iteration))
genA2B.save_weights(args.model_save_dir + '/genA2B_' +
str(iteration))
if iteration % args.test_interval == 0:
input_images = []
output_images = []
for _ in range(args.test_size):
test_full = next(test_dataset)
split_list_test = pyramid.split(test_full, args.level)
test_low = split_list_test[-1]
generated_low = genA2B(test_low, training=False)
generated_full = merge_image(generated_low, split_list_test)
input_images.append(test_full)
output_images.append(generated_full)
generate_images(input_images[0], input_images[1], output_images[1],
output_images[0], args.save_dir_test, iteration)
def train(train_datasetA, train_datasetB, args):
start = time.time()
for iteration in range(args.start_iter, args.max_iterations):
with tf.GradientTape() as genA2B_tape, tf.GradientTape() as genB2A_tape, \
tf.GradientTape() as discA_tape, tf.GradientTape() as discB_tape:
try:
trainA_full = next(train_datasetA)
trainB_full = next(train_datasetB)
except tf.errors.OutOfRangeError:
print("Error, run out of data")
break
split_list_A = pyramid.split(trainA_full, args.level)
split_list_B = pyramid.split(trainB_full, args.level)
trainA = split_list_A[-1]
trainB = split_list_B[-1]
genA2B_output = genA2B(trainA, training=True)
genB2A_output = genB2A(trainB, training=True)
genA2B_output_full = merge_image(genA2B_output, split_list_A)
genB2A_output_full = merge_image(genB2A_output, split_list_B)
reconstructed_A_low = genB2A(genA2B_output, training=True)
reconstructed_B_low = genA2B(genB2A_output, training=True)
reconstructed_A_full = merge_image(reconstructed_A_low,
split_list_A)
reconstructed_B_full = merge_image(reconstructed_B_low,
split_list_B)
discA_real_output = discA(trainA_full)
discB_real_output = discB(trainB_full)
discA_fake_output = discA(genB2A_output_full)
discB_fake_output = discB(genA2B_output_full)
discA_loss = discriminator_losssss('lsgan', discA_real_output,
discA_fake_output)
discB_loss = discriminator_losssss('lsgan', discB_real_output,
discB_fake_output)
color_loss_A2B = color_losses(genA2B_output_full, trainB_full,
args.batch_size)
color_loss_B2A = color_losses(genB2A_output_full, trainA_full,
args.batch_size)
# content_loss_A2B = content_losses(genA2B_output_full, trainA_full, args.batch_size)
# content_loss_B2A = content_losses(genB2A_output_full, trainB_full, args.batch_size)
# gp_A = gradient_penalty(trainA_full, genB2A_output_full, discA)
# gp_B = gradient_penalty(trainB_full, genA2B_output_full, discB)
generatorA2B_loss = generator_losssss('lsgan', discB_fake_output)
cycleA2B_loss = cycle_consistency_loss(trainA_full, trainB_full,
reconstructed_A_full,
reconstructed_B_full)
generatorB2A_loss = generator_losssss('lsgan', discA_fake_output)
cycleB2A_loss = cycle_consistency_loss(trainA_full, trainB_full,
reconstructed_A_full,
reconstructed_B_full)
genA2B_loss = args.dis_lambda * generatorA2B_loss + args.cyc_lambda * cycleA2B_loss + args.color_lambda * color_loss_A2B
genB2A_loss = args.dis_lambda * generatorB2A_loss + args.cyc_lambda * cycleB2A_loss + args.color_lambda * color_loss_B2A
genA2B_gradients = genA2B_tape.gradient(genA2B_loss,
genA2B.trainable_variables)
genB2A_gradients = genB2A_tape.gradient(genB2A_loss,
genB2A.trainable_variables)
discA_gradients = discA_tape.gradient(discA_loss,
discA.trainable_variables)
discA_gradients, _ = tf.clip_by_global_norm(discA_gradients, 15)
discB_gradients = discB_tape.gradient(discB_loss,
discB.trainable_variables)
discB_gradients, _ = tf.clip_by_global_norm(discB_gradients, 15)
genA2B_optimizer.apply_gradients(
zip(genA2B_gradients, genA2B.trainable_variables))
genB2A_optimizer.apply_gradients(
zip(genB2A_gradients, genB2A.trainable_variables))
discA_optimizer.apply_gradients(
zip(discA_gradients, discA.trainable_variables))
discB_optimizer.apply_gradients(
zip(discB_gradients, discB.trainable_variables))
if iteration % args.log_interval == 0:
print(
'Training ' + procname +
', Iteration: {}th, Duration: {:.4f}, LOSSES: cycle: {:.4f}, generator: {:.4f}, disc: {:.4f}, color: {:.4f}'
.format(iteration,
time.time() -
start, cycleA2B_loss, generatorA2B_loss.numpy(),
discA_loss.numpy(), color_loss_A2B))
start = time.time()
if iteration % args.save_interval == 0:
generate_images(trainA_full, trainB_full, genA2B_output_full,
genB2A_output_full, save_dir_train, iteration)
# print('Time taken for iteration {} is {} sec'.format(iteration + 1, time.time() - start))
if not os.path.exists(model_save_dir):
os.mkdir(model_save_dir)
# checkpoint.save(file_prefix = checkpoint_prefix)
discA.save_weights(model_save_dir + '/discA_' + str(iteration))
discB.save_weights(model_save_dir + '/discB_' + str(iteration))
genA2B.save_weights(model_save_dir + '/genA2B_' + str(iteration))
genB2A.save_weights(model_save_dir + '/genB2A_' + str(iteration))
if iteration % args.test_interval == 0:
input_images = []
output_images = []
test_full = next(test_datasetA)
start_test = time.time()
split_list_test = pyramid.split(test_full, args.level)
test_low = split_list_test[-1]
generated_low = genA2B(test_low, training=False)
generated_full = merge_image(generated_low, split_list_test)
test_time_1 = time.time() - start_test
input_images.append(test_full)
output_images.append(generated_full)
test_full = next(test_datasetB)
start_test = time.time()
split_list_test = pyramid.split(test_full, args.level)
test_low = split_list_test[-1]
generated_low = genB2A(test_low, training=False)
generated_full = merge_image(generated_low, split_list_test)
test_time_2 = time.time() - start_test
input_images.append(test_full)
output_images.append(generated_full)
generate_images(input_images[0], input_images[1], output_images[0],
output_images[1], save_dir_test, iteration)
if iteration % 1000 == 0:
print('test time: ', test_time_1)
print('test time: ', test_time_2)
new_shape = [image_size, image_size]
img = tf.image.resize(img, new_shape)
img = img[tf.newaxis, :]
return img
def imshow(image, title=None):
if len(image.shape) > 3:
image = tf.squeeze(image, axis=0)
content_image = load_img(content_path)
style_image = load_img(style_path)
style_image = pyramid.split(style_image, level)[-1]
vgg = tf.keras.applications.VGG19(include_top=False, weights='imagenet')
content_layers = ['block5_conv2']
style_layers = [
'block1_conv1', 'block2_conv1', 'block3_conv1', 'block4_conv1',
'block5_conv1'
]
num_content_layers = len(content_layers)
num_style_layers = len(style_layers)
def vgg_layers(layer_names):
def train(train_datasetA, train_datasetB, args, lsgan=True):
procname = os.path.basename(args.model_save_dir)
setproctitle.setproctitle('train_{}'.format(procname))
start = time.time()
for iteration in range(args.start_iter, args.max_iterations):
with tf.GradientTape() as genA2B_tape, tf.GradientTape() as genB2A_tape, \
tf.GradientTape() as discA_tape, tf.GradientTape() as discB_tape:
try:
trainA_full = next(train_datasetA)
trainB_full = next(train_datasetB)
except tf.errors.OutOfRangeError:
print("Error, run out of data")
break
split_list_A = pyramid.split(trainA_full, args.level)
split_list_B = pyramid.split(trainB_full, args.level)
trainA = split_list_A[-1]
trainB = split_list_B[-1]
genA2B_output = genA2B(trainA, training=True)
genB2A_output = genB2A(trainB, training=True)
genA2B_output_full = merge_image(genA2B_output, split_list_A)
genB2A_output_full = merge_image(genB2A_output, split_list_B)
split_output_A = pyramid.split(genA2B_output_full, args.level)
split_output_B = pyramid.split(genB2A_output_full, args.level)
reconstructed_A_low = genB2A(split_output_A[-1], training=True)
reconstructed_B_low = genA2B(split_output_B[-1], training=True)
# reconstructed_A = genB2A(genA2B_output, training=True)
# reconstructed_B = genA2B(genB2A_output, training=True)
# genA2B_output = genA2B(trainA, training=True)
# genB2A_output = genB2A(trainB, training=True)
reconstructedA_full = merge_image(reconstructed_A_low,
split_output_A)
reconstructedB_full = merge_image(reconstructed_B_low,
split_output_B)
discA_real_output = discA(trainA_full, training=True)
discB_real_output = discB(trainB_full, training=True)
discA_fake_output = discA(genB2A_output_full, training=True)
discB_fake_output = discB(genA2B_output_full, training=True)
# Use history buffer of 50 for disc loss
discA_loss = discriminator_loss(discA_real_output,
discA_fake_output,
lsgan=lsgan)
discB_loss = discriminator_loss(discB_real_output,
discB_fake_output,
lsgan=lsgan)
generatorA2B_loss = generator_loss(discB_fake_output, lsgan=lsgan)
cycleA2B_loss = cycle_consistency_loss(trainA_full,
trainB_full,
reconstructedA_full,
reconstructedB_full,
cyc_lambda=args.cyc_lambda)
generatorB2A_loss = generator_loss(discA_fake_output, lsgan=lsgan)
cycleB2A_loss = cycle_consistency_loss(trainA_full,
trainB_full,
reconstructedA_full,
reconstructedB_full,
cyc_lambda=args.cyc_lambda)
genA2B_loss = args.dis_lambda * generatorA2B_loss + cycleA2B_loss
genB2A_loss = args.dis_lambda * generatorB2A_loss + cycleB2A_loss
genA2B_gradients = genA2B_tape.gradient(genA2B_loss,
genA2B.trainable_variables)
genB2A_gradients = genB2A_tape.gradient(genB2A_loss,
genB2A.trainable_variables)
discA_gradients = discA_tape.gradient(discA_loss,
discA.trainable_variables)
discB_gradients = discB_tape.gradient(discB_loss,
discB.trainable_variables)
genA2B_optimizer.apply_gradients(
zip(genA2B_gradients, genA2B.trainable_variables))
genB2A_optimizer.apply_gradients(
zip(genB2A_gradients, genB2A.trainable_variables))
discA_optimizer.apply_gradients(
zip(discA_gradients, discA.trainable_variables))
discB_optimizer.apply_gradients(
zip(discB_gradients, discB.trainable_variables))
if iteration % args.log_interval == 0:
# print('Training Iteration: {}th, LOSSES: D_A: {}, D_B: {}, G_A: {}, G_B: {}, cycle: {}'.format(
# iteration, discA_loss, discB_loss, generatorA2B_loss, generatorB2A_loss, cycleA2B_loss))
# print('Training Iteration: {}th, LOSSES: cycle: {:.4f}'.format(
# iteration, cycleA2B_loss))
print(
'Training ' + procname +
', Iteration: {}th, time: {:.4f}, LOSSES: cycle: {:.4f}, discir: {:.4f}'
.format(iteration,
time.time() - start, cycleA2B_loss, generatorA2B_loss))
start = time.time()
if iteration % args.save_interval == 0:
generate_images(trainA_full, trainB_full, genB2A_output_full,
genA2B_output_full, args.save_dir_train, iteration)
# print('Time taken for iteration {} is {} sec'.format(iteration + 1, time.time() - start))
if not os.path.exists(args.model_save_dir):
os.mkdir(args.model_save_dir)
# checkpoint.save(file_prefix = checkpoint_prefix)
discA.save_weights(args.model_save_dir + '/discA_' +
str(iteration))
discB.save_weights(args.model_save_dir + '/discB_' +
str(iteration))
genA2B.save_weights(args.model_save_dir + '/genA2B_' +
str(iteration))
genB2A.save_weights(args.model_save_dir + '/genB2A_' +
str(iteration))
if iteration % args.test_interval == 0:
input_images = []
output_images = []
for _ in range(args.test_size):
test_full = next(test_dataset)
split_list_test = pyramid.split(test_full, args.level)
test_low = split_list_test[-1]
generated_low = genA2B(test_low, training=False)
generated_full = merge_image(generated_low, split_list_test)
input_images.append(test_full)
output_images.append(generated_full)
generate_images(input_images[0], input_images[1], output_images[0],
output_images[1], args.save_dir_test, iteration)