def train_step(lr, hr, generator, discriminator, content):
with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
## re-scale
## lr: 0 ~ 1
## hr: -1 ~ 1
lr = tf.cast(lr, tf.float32)
hr = tf.cast(hr, tf.float32)
lr = lr / 255
hr = hr / 127.5 - 1
sr = generator(lr, training=True)
sr_output = discriminator(sr, training=True)
hr_output = discriminator(hr, training=True)
disc_loss = discriminator_loss(sr_output, hr_output)
mse_loss = mse_based_loss(sr, hr)
gen_loss = generator_loss(sr_output)
cont_loss = content_loss(content, sr, hr)
perceptual_loss = mse_loss + cont_loss + 1e-3 * gen_loss
gradients_of_generator = gen_tape.gradient(
perceptual_loss, generator.trainable_variables)
gradients_of_discriminator = disc_tape.gradient(
disc_loss, discriminator.trainable_variables)
generator_optimizer.apply_gradients(
zip(gradients_of_generator, generator.trainable_variables))
discriminator_optimizer.apply_gradients(
zip(gradients_of_discriminator, discriminator.trainable_variables))
return perceptual_loss, disc_loss
def dis_step(input_image, target):
with tf.GradientTape() as disc_tape:
gen_output = generator(input_image, training=True)
disc_real_output = discriminator([input_image, target], training=True)
disc_generated_output = discriminator([input_image, gen_output], training=True)
disc_loss,real_loss,generated_loss = discriminator_loss(disc_real_output, disc_generated_output)
discriminator_gradients = disc_tape.gradient(disc_loss,
discriminator.trainable_variables)
discriminator_optimizer.apply_gradients(zip(discriminator_gradients,
discriminator.trainable_variables))
return disc_loss,real_loss,generated_loss
def train_step(molecules_A, molecules_X) :
"""
First generates molecules with the generator, then pass a batch of data and the generated molecules through the discriminator, then applies backpropagation
"""
z = np.random.randn(molecules_A.shape[0], 32)
with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape :
generated_A, generated_X = generator(z)
real_logits = discriminator(molecules_A, molecules_X)
fake_logits = discriminator(generated_A[0], generated_X[0])
# backpropagation
# gradient penalty : WGAN, discriminator loss
with gen_tape.stop_recording() :
with disc_tape.stop_recording() :
# cf. equation (3) in the paper : penalty on the gradient of the discriminator wrt a linear combination of real and generated data
epsilon_adj = tf.random.uniform(tf.shape(molecules_A), 0.0, 1.0, dtype=molecules_A.dtype)
epsilon_features = tf.random.uniform(tf.shape(molecules_X), 0.0, 1.0, dtype=generated_X.dtype)
with tf.GradientTape() as penalty_tape :
m1 = epsilon_adj * molecules_A
m2 = (1 - epsilon_adj)*generated_A[0]
x_hat_adj = m1 + m2
x_hat_features = epsilon_features * molecules_X + (1 - epsilon_features) * generated_X[0]
penalty_tape.watch([x_hat_adj, x_hat_features])
disc_penalty = discriminator(x_hat_adj, x_hat_features)
# get the gradient, again eq (3) in the paper
grad_adj = penalty_tape.gradient(disc_penalty, [x_hat_adj, x_hat_features])
disc_loss = loss.discriminator_loss(real_logits, fake_logits, grad_adj)
gen_loss = loss.generator_loss(fake_logits)
gradients_of_generator = gen_tape.gradient(gen_loss, generator.variables)
gradients_of_discriminator = disc_tape.gradient(disc_loss, discriminator.variables)
disciminator_optimizer.apply_gradients(zip(gradients_of_discriminator, discriminator.variables))
generator_optimizer.apply_gradients(zip(gradients_of_generator, generator.variables))
gen_train_loss(gen_loss)
disc_train_loss(disc_loss)
return gen_train_loss, disc_train_loss
def train_step(input_image, target,LAMBDA):
with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
gen_output = generator(input_image, training=True)
disc_real_output = discriminator([input_image, target], training=True)
disc_generated_output = discriminator([input_image, gen_output], training=True)
gen_total_loss, gen_gan_loss, gen_l1_loss = generator_loss(disc_generated_output, gen_output, target,LAMBDA)
disc_loss,real_loss,generated_loss = discriminator_loss(disc_real_output, disc_generated_output)
generator_gradients = gen_tape.gradient(gen_total_loss,
generator.trainable_variables)
discriminator_gradients = disc_tape.gradient(disc_loss,
discriminator.trainable_variables)
generator_optimizer.apply_gradients(zip(generator_gradients,
generator.trainable_variables))
discriminator_optimizer.apply_gradients(zip(discriminator_gradients,
discriminator.trainable_variables))
return gen_total_loss, gen_gan_loss, gen_l1_loss,disc_loss,real_loss,generated_loss
fake_tex = model_d(fake_img)
# Update Discriminator network
if iteration % opt.disc_freq == 0:
# calculate gradient penalty
if opt.wgan:
rand = torch.rand(1).item()
sample = rand * disc_img + (1 - rand) * fake_img
gp_tex = model_d(sample)
gradient = torch.autograd.grad(gp_tex.mean(), sample, create_graph=True)[0]
grad_pen = 10 * (gradient.norm() - 1) ** 2
else:
grad_pen = None
# update discriminator
model_d.zero_grad()
d_tex_loss = loss.discriminator_loss(real_tex, fake_tex, wasserstein=opt.wgan, grad_penalties=grad_pen)
d_tex_loss.backward(retain_graph=True)
optimizer_d.step()
# save data to tensorboard
if opt.saving:
writer.add_scalar('loss/d_tex_loss', d_tex_loss, iteration)
if opt.wgan:
writer.add_scalar('disc_score/gradient_penalty', grad_pen.mean().data.item(), iteration)
# Update Generator network
if iteration % opt.gen_freq == 0:
# update discriminator
model_g.zero_grad()
if opt.generator == 'DSGAN':
g_loss = g_loss_module(fake_tex, fake_img, input_img)
elif opt.generator == 'DeResnet':
def main(args):
os.makedirs(args.log_dir, exist_ok=True)
# create models
G_1 = Generator_lr(in_channels=3)
G_2 = Generator_lr(in_channels=3)
D_1 = Discriminator_lr(in_channels=3, in_h=16, in_w=16)
SR = EDSR(n_colors=3)
G_3 = Generator_sr(in_channels=3)
D_2 = Discriminator_sr(in_channels=3, in_h=64, in_w=64)
for model in [G_1, G_2, D_1, SR, G_3, D_2]:
model.cuda()
model.train()
# tensorboard
writer = SummaryWriter(log_dir=args.log_dir)
# create optimizors
optim = {
'G_1':
torch.optim.Adam(params=filter(lambda p: p.requires_grad,
G_1.parameters()),
lr=args.lr * 5),
'G_2':
torch.optim.Adam(params=filter(lambda p: p.requires_grad,
G_2.parameters()),
lr=args.lr * 5),
'D_1':
torch.optim.Adam(params=filter(lambda p: p.requires_grad,
D_1.parameters()),
lr=args.lr),
'SR':
torch.optim.Adam(params=filter(lambda p: p.requires_grad,
SR.parameters()),
lr=args.lr * 5),
'G_3':
torch.optim.Adam(params=filter(lambda p: p.requires_grad,
G_3.parameters()),
lr=args.lr),
'D_2':
torch.optim.Adam(params=filter(lambda p: p.requires_grad,
D_2.parameters()),
lr=args.lr)
}
for key in optim.keys():
optim[key].zero_grad()
# get dataloader
train_dataset = DIV2KDataset(root=args.data_path)
trainloader = DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=3)
print('-' * 20)
print('Start training')
print('-' * 20)
iter_index = 0
for epoch in range(args.epochs):
G_1.train()
SR.train()
start = timeit.default_timer()
for _, batch in enumerate(trainloader):
iter_index += 1
image, label_hr, label_lr = batch
image = image.cuda()
label_hr = label_hr.cuda()
label_lr = label_lr.cuda()
'''loss for lr GAN'''
'''update G_1 and G_2'''
for key in optim.keys():
optim[key].zero_grad()
# D loss for D_1
image_clean = G_1(image)
loss_D1 = discriminator_loss(discriminator=D_1,
fake=image_clean,
real=label_lr)
loss_D1.backward()
optim['D_1'].step()
# GD loss for G_1
loss_G1 = generator_discriminator_loss(generator=G_1,
discriminator=D_1,
input=image)
loss_G1.backward()
# cycle loss for G_1 and G_2
loss_cycle = 10 * cycle_loss(G_1, G_2, image)
loss_cycle.backward()
# idt loss for G_1
loss_idt = 5 * identity_loss(clean_image=label_lr, generator=G_1)
loss_idt.backward()
# tvloss for G_1
loss_tv = 0.5 * tvloss(input=image, generator=G_1)
loss_tv.backward()
# optimize G_1 and G_2
optim['G_1'].step()
optim['G_2'].step()
if iter_index % 100 == 0:
print(
'iter {}: LR: loss_D1={}, loss_GD={}, loss_cycle={}, loss_idt={}, loss_tv={}'
.format(iter_index, loss_D1.item(), loss_G1.item(),
loss_cycle.item(), loss_idt.item(),
loss_tv.item()))
writer.add_scalar('LR/loss_D1', loss_D1.item(),
iter_index // 100)
writer.add_scalar('LR/loss_GD', loss_G1.item(),
iter_index // 100)
writer.add_scalar('LR/loss_cycle', loss_cycle.item(),
iter_index // 100)
writer.add_scalar('LR/loss_idt', loss_idt.item(),
iter_index // 100)
writer.add_scalar('LR/loss_tv', loss_tv.item(),
iter_index // 100)
writer.add_image('LR/origin', image[0], iter_index // 100)
writer.add_image('LR/denoise',
G_1(image)[0], iter_index // 100)
'''loss for sr GAN'''
'''update G_1, SR and G_3'''
for key in optim.keys():
optim[key].zero_grad()
image_clean = G_1(image).detach()
# D loss for D_2
image_sr = SR(image_clean)
loss_D2 = discriminator_loss(discriminator=D_2,
fake=image_sr,
real=label_hr)
loss_D2.backward()
optim['D_2'].step()
# GD loss for SR
loss_SR = generator_discriminator_loss(generator=SR,
discriminator=D_2,
input=image_clean)
loss_SR.backward()
# cycle loss for SR and G_3
loss_cycle = 10 * cycle_loss(SR, G_3, image_clean)
loss_cycle.backward()
# idt loss for SR
loss_idt = 5 * identity_loss_sr(
clean_image_lr=label_lr, clean_image_hr=label_hr, generator=SR)
loss_idt.backward()
# tvloss for SR
loss_tv = 0.5 * tvloss(input=image_clean, generator=SR)
loss_tv.backward()
# optimize G_1, SR and G_3
optim['G_1'].step()
optim['SR'].step()
optim['G_3'].step()
if iter_index % 100 == 0:
print(
' SR: loss_D2={}, loss_SR={}, loss_cycle={}, loss_idt={}, loss_tv={}'
.format(loss_D2.item(), loss_SR.item(), loss_cycle.item(),
loss_idt.item(), loss_tv.item()))
writer.add_scalar('SR/loss_D2', loss_D2.item(),
iter_index // 100)
writer.add_scalar('SR/loss_SR', loss_SR.item(),
iter_index // 100)
writer.add_scalar('SR/loss_cycle', loss_cycle.item(),
iter_index // 100)
writer.add_scalar('SR/loss_idt', loss_idt.item(),
iter_index // 100)
writer.add_scalar('SR/loss_tv', loss_tv.item(),
iter_index // 100)
writer.add_image('SR/origin', image[0], iter_index // 100)
writer.add_image('SR/clean_image',
G_1(image)[0], iter_index // 100)
writer.add_image('SR/SR', SR(G_1(image))[0], iter_index // 100)
writer.flush()
end = timeit.default_timer()
print('epoch {}, using {} seconds'.format(epoch, end - start))
G_1.eval()
SR.eval()
image = Image.open('/data/data/DIV2K/unsupervised/lr/0001x4d.png')
sr_image = resolv_sr(G_1, SR, image)
# image_tensor = torchvision.transforms.functional.to_tensor(image).unsqueeze(0).cuda()
# sr_image_tensor = SR(G_1(image_tensor).detach())
# sr_image = torchvision.transforms.functional.to_pil_image(sr_image_tensor[0].cpu())
sr_image.save(
os.path.join(args.log_dir, '0001x4d_sr_{}.png'.format(str(epoch))))
torch.save(G_1.state_dict(),
os.path.join(args.log_dir, 'ep-' + str(epoch) + '_G_1.pkl'))
torch.save(G_2.state_dict(),
os.path.join(args.log_dir, 'ep-' + str(epoch) + '_G_2.pkl'))
torch.save(D_1.state_dict(),
os.path.join(args.log_dir, 'ep-' + str(epoch) + '_D_1.pkl'))
torch.save(SR.state_dict(),
os.path.join(args.log_dir, 'ep-' + str(epoch) + '_SR.pkl'))
torch.save(G_3.state_dict(),
os.path.join(args.log_dir, 'ep-' + str(epoch) + '_G_3.pkl'))
torch.save(D_2.state_dict(),
os.path.join(args.log_dir, 'ep-' + str(epoch) + '_D_2.pkl'))
writer.close()
print('Training done.')
torch.save(G_1.state_dict(),
os.path.join(args.log_dir, 'final_weights_G_1.pkl'))
torch.save(G_2.state_dict(),
os.path.join(args.log_dir, 'final_weights_G_2.pkl'))
torch.save(D_1.state_dict(),
os.path.join(args.log_dir, 'final_weights_D_1.pkl'))
torch.save(SR.state_dict(),
os.path.join(args.log_dir, 'final_weights_SR.pkl'))
torch.save(G_3.state_dict(),
os.path.join(args.log_dir, 'final_weights_G_3.pkl'))
torch.save(D_2.state_dict(),
os.path.join(args.log_dir, 'final_weights_D_2.pkl'))
image = Image.open('/data/data/DIV2K/unsupervised/lr/0001x4d.png')
image.save(os.path.join(args.log_dir, '0001x4d.png'))
sr_image = resolv_sr(G_1, SR, image)
# image_tensor = torchvision.transforms.functional.to_tensor(image).unsqueeze(0).cuda()
# sr_image_tensor = SR(G_1(image_tensor))
# sr_image = torchvision.transforms.functional.to_pil_image(sr_image_tensor[0].cpu())
sr_image.save(os.path.join(args.log_dir, '0001x4d_sr.png'))
def trainstep(real_human, real_anime, big_anime):
with tf.GradientTape(persistent=True) as tape:
fake_anime = generator_to_anime(real_human, training=True)
cycled_human = generator_to_human(fake_anime, training=True)
fake_human = generator_to_human(real_anime, training=True)
cycled_anime = generator_to_anime(fake_human, training=True)
# same_human and same_anime are used for identity loss.
same_human = generator_to_human(real_human, training=True)
same_anime = generator_to_anime(real_anime, training=True)
disc_real_human = discriminator_human(real_human, training=True)
disc_real_anime = discriminator_anime(real_anime, training=True)
disc_fake_human = discriminator_human(fake_human, training=True)
disc_fake_anime = discriminator_anime(fake_anime, training=True)
# calculate the loss
gen_anime_loss = generator_loss(disc_fake_anime)
gen_human_loss = generator_loss(disc_fake_human)
total_cycle_loss = cycle_loss(real_human,
cycled_human) + cycle_loss(
real_anime, cycled_anime)
# Total generator loss = adversarial loss + cycle loss
total_gen_anime_loss = (gen_anime_loss + total_cycle_loss +
identity_loss(real_anime, same_anime))
total_gen_anime_loss = generator_to_anime_optimizer.get_scaled_loss(
total_gen_anime_loss)
total_gen_human_loss = (gen_human_loss + total_cycle_loss +
identity_loss(real_human, same_human))
total_gen_human_loss = generator_to_human_optimizer.get_scaled_loss(
total_gen_human_loss)
disc_human_loss = discriminator_loss(disc_real_human,
disc_fake_human)
disc_human_loss = discriminator_human_optimizer.get_scaled_loss(
disc_human_loss)
disc_anime_loss = discriminator_loss(disc_real_anime,
disc_fake_anime)
disc_anime_loss = discriminator_anime_optimizer.get_scaled_loss(
disc_anime_loss)
# My part
fake_anime_upscale = generator_anime_upscale(fake_anime)
cycle_anime_upscale = generator_anime_upscale(real_anime)
same_anime_upscale = generator_anime_upscale(same_anime)
disc_fake_upscale = discriminator_anime_upscale(fake_anime_upscale)
disc_cycle_upscale = discriminator_anime_upscale(
cycle_anime_upscale)
disc_same_upscale = discriminator_anime_upscale(same_anime_upscale)
disc_real_big = discriminator_anime_upscale(big_anime)
gen_upscale_loss = (
generator_loss(disc_fake_upscale) * 3 +
generator_loss(disc_cycle_upscale) +
generator_loss(disc_same_upscale)
# + mse_loss(big_anime, cycle_anime_upscale)
+ identity_loss(big_anime, cycle_anime_upscale) * 0.5 +
identity_loss(big_anime, same_anime_upscale) * 0.5)
gen_upscale_loss = generator_anime_upscale_optimizer.get_scaled_loss(
gen_upscale_loss)
disc_upscale_loss = discriminator_upscale_loss(
disc_real_big, disc_fake_upscale, disc_cycle_upscale,
disc_same_upscale)
disc_upscale_loss = discriminator_anime_upscale_optimizer.get_scaled_loss(
disc_upscale_loss)
# Calculate the gradients for generator and discriminator
generator_to_anime_gradients = tape.gradient(
total_gen_anime_loss, generator_to_anime.trainable_variables)
generator_to_human_gradients = tape.gradient(
total_gen_human_loss, generator_to_human.trainable_variables)
discriminator_human_gradients = tape.gradient(
disc_human_loss, discriminator_human.trainable_variables)
discriminator_anime_gradients = tape.gradient(
disc_anime_loss, discriminator_anime.trainable_variables)
generator_upscale_gradients = tape.gradient(
gen_upscale_loss, generator_anime_upscale.trainable_variables)
discriminator_upscale_gradients = tape.gradient(
disc_upscale_loss, discriminator_anime_upscale.trainable_variables)
# Apply the gradients to the optimizer
generator_to_anime_gradients = generator_to_anime_optimizer.get_unscaled_gradients(
generator_to_anime_gradients)
generator_to_anime_optimizer.apply_gradients(
zip(generator_to_anime_gradients,
generator_to_anime.trainable_variables))
generator_to_human_gradients = generator_to_human_optimizer.get_unscaled_gradients(
generator_to_human_gradients)
generator_to_human_optimizer.apply_gradients(
zip(generator_to_human_gradients,
generator_to_human.trainable_variables))
discriminator_human_gradients = discriminator_human_optimizer.get_unscaled_gradients(
discriminator_human_gradients)
discriminator_human_optimizer.apply_gradients(
zip(discriminator_human_gradients,
discriminator_human.trainable_variables))
discriminator_anime_gradients = discriminator_anime_optimizer.get_unscaled_gradients(
discriminator_anime_gradients)
discriminator_anime_optimizer.apply_gradients(
zip(discriminator_anime_gradients,
discriminator_anime.trainable_variables))
generator_upscale_gradients = generator_anime_upscale_optimizer.get_unscaled_gradients(
generator_upscale_gradients)
generator_anime_upscale_optimizer.apply_gradients(
zip(generator_upscale_gradients,
generator_anime_upscale.trainable_variables))
discriminator_upscale_gradients = discriminator_anime_upscale_optimizer.get_unscaled_gradients(
discriminator_upscale_gradients)
discriminator_anime_upscale_optimizer.apply_gradients(
zip(
discriminator_upscale_gradients,
discriminator_anime_upscale.trainable_variables,
))
return (
real_human,
real_anime,
fake_anime,
cycled_human,
fake_human,
cycled_anime,
same_human,
same_anime,
fake_anime_upscale,
same_anime_upscale,
gen_anime_loss,
gen_human_loss,
disc_human_loss,
disc_anime_loss,
total_gen_anime_loss,
total_gen_human_loss,
gen_upscale_loss,
disc_upscale_loss,
)
def trainstep(real_human, real_anime):
with tf.GradientTape(persistent=True) as tape:
fake_anime = generator_to_anime(real_human, training=True)
cycled_human = generator_to_human(fake_anime, training=True)
fake_human = generator_to_human(real_anime, training=True)
cycled_anime = generator_to_anime(fake_human, training=True)
# same_human and same_anime are used for identity loss.
same_human = generator_to_human(real_human, training=True)
same_anime = generator_to_anime(real_anime, training=True)
disc_real_human = discriminator_x(real_human, training=True)
disc_real_anime = discriminator_y(real_anime, training=True)
disc_fake_human = discriminator_x(fake_human, training=True)
disc_fake_anime = discriminator_y(fake_anime, training=True)
# calculate the loss
gen_anime_loss = generator_loss(disc_fake_anime)
gen_human_loss = generator_loss(disc_fake_human)
total_cycle_loss = cycle_loss(real_human,
cycled_human) + cycle_loss(
real_anime, cycled_anime)
# Total generator loss = adversarial loss + cycle loss
total_gen_anime_loss = gen_anime_loss + total_cycle_loss + identity_loss(
real_anime, same_anime)
total_gen_human_loss = gen_human_loss + total_cycle_loss + identity_loss(
real_human, same_human)
disc_x_loss = discriminator_loss(disc_real_human, disc_fake_human)
disc_y_loss = discriminator_loss(disc_real_anime, disc_fake_anime)
# Calculate the gradients for generator and discriminator
generator_to_anime_gradients = tape.gradient(
total_gen_anime_loss, generator_to_anime.trainable_variables)
generator_to_human_gradients = tape.gradient(
total_gen_human_loss, generator_to_human.trainable_variables)
discriminator_x_gradients = tape.gradient(
disc_x_loss, discriminator_x.trainable_variables)
discriminator_y_gradients = tape.gradient(
disc_y_loss, discriminator_y.trainable_variables)
# Apply the gradients to the optimizer
generator_to_anime_optimizer.apply_gradients(
zip(generator_to_anime_gradients,
generator_to_anime.trainable_variables))
generator_to_human_optimizer.apply_gradients(
zip(generator_to_human_gradients,
generator_to_human.trainable_variables))
discriminator_x_optimizer.apply_gradients(
zip(discriminator_x_gradients,
discriminator_x.trainable_variables))
discriminator_y_optimizer.apply_gradients(
zip(discriminator_y_gradients,
discriminator_y.trainable_variables))
return fake_anime, cycled_human, fake_human, cycled_anime , same_human , same_anime, \
gen_anime_loss, gen_human_loss, disc_x_loss, disc_y_loss, total_gen_anime_loss, total_gen_human_loss
def trainstep(real_human, real_anime, big_anime):
with tf.GradientTape(persistent=True) as tape:
ones = tf.ones_like(real_human)
neg_ones = tf.ones_like(real_human) * -1
def get_domain_anime(img):
return tf.concat([img, ones], 3)
def get_domain_human(img):
return tf.concat([img, neg_ones], 3)
fake_anime = generator(get_domain_anime(real_human), training=True)
cycled_human = generator(get_domain_human(fake_anime), training=True)
fake_human = generator(get_domain_human(real_anime), training=True)
cycled_anime = generator(get_domain_anime(fake_human), training=True)
# same_human and same_anime are used for identity loss.
same_human = generator(get_domain_human(real_human), training=True)
same_anime = generator(get_domain_anime(real_anime), training=True)
disc_real_human, label_real_human = discriminator(real_human, training=True)
disc_real_anime, label_real_anime = discriminator(real_anime, training=True)
disc_fake_human, label_fake_human = discriminator(fake_human, training=True)
disc_fake_anime, label_fake_anime = discriminator(fake_anime, training=True)
_, label_cycled_human = discriminator(cycled_human, training=True)
_, label_cycled_anime = discriminator(cycled_anime, training=True)
_, label_same_human = discriminator(same_human, training=True)
_, label_same_anime = discriminator(same_anime, training=True)
# calculate the loss
gen_anime_loss = generator_loss(disc_fake_anime)
gen_human_loss = generator_loss(disc_fake_human)
total_cycle_loss = cycle_loss(real_human, cycled_human) + cycle_loss(
real_anime, cycled_anime
)
gen_class_loss = (
discriminator_loss(label_fake_human, label_fake_anime)
+ discriminator_loss(label_cycled_human, label_cycled_anime)
+ discriminator_loss(label_same_human, label_same_anime)
)
# Total generator loss = adversarial loss + cycle loss
total_gen_loss = (
gen_anime_loss
+ gen_human_loss
+ gen_class_loss
+ total_cycle_loss * 0.1
+ identity_loss(real_anime, same_anime)
+ identity_loss(real_human, same_human)
)
tf.print("gen_anime_loss",gen_anime_loss)
tf.print("gen_human_loss",gen_human_loss)
tf.print("gen_class_loss",gen_class_loss)
tf.print("total_cycle_loss",total_cycle_loss)
tf.print("identity_loss(real_anime, same_anime)",identity_loss(real_anime, same_anime))
tf.print("identity_loss(real_human, same_human)",identity_loss(real_human, same_human))
scaled_total_gen_anime_loss = generator_optimizer.get_scaled_loss(
total_gen_loss
)
disc_human_loss = discriminator_loss(disc_real_human, disc_fake_human)
disc_anime_loss = discriminator_loss(disc_real_anime, disc_fake_anime)
# disc_gp_anime = gradient_penalty_star(partial(discriminator, training=True), real_anime,fake_anime )
# disc_gp_human = gradient_penalty_star(partial(discriminator, training=True), real_human,fake_human )
disc_loss = disc_human_loss + disc_anime_loss + discriminator_loss(label_real_human,label_real_anime)
# +disc_gp_anime+disc_gp_human
scaled_disc_loss = discriminator_optimizer.get_scaled_loss(
disc_loss
)
# Calculate the gradients for generator and discriminator
generator_gradients =generator_optimizer.get_unscaled_gradients( tape.gradient(
scaled_total_gen_anime_loss, generator.trainable_variables
))
discriminator_gradients = discriminator_optimizer.get_unscaled_gradients( tape.gradient(
scaled_disc_loss, discriminator.trainable_variables
))
generator_optimizer.apply_gradients(
zip(generator_gradients, generator.trainable_variables)
)
discriminator_optimizer.apply_gradients(
zip(discriminator_gradients, discriminator.trainable_variables)
)
with tf.GradientTape(persistent=True) as tape:
real_anime_up = up_G(real_anime)
fake_anime_up = up_G(fake_anime)
dis_fake_anime_up = up_D(fake_anime_up)
dis_real_anime_up = up_D(real_anime_up)
dis_ori_anime = up_D(big_anime)
gen_up_loss = generator_loss(fake_anime_up) + generator_loss(dis_real_anime_up)*0.1
dis_up_loss = discriminator_loss(dis_ori_anime,dis_fake_anime_up)+discriminator_loss(dis_ori_anime,dis_real_anime_up)*0.1
scaled_gen_up_loss = up_G_optim.get_scaled_loss(gen_up_loss)
scaled_disc_loss = up_D_optim.get_scaled_loss(dis_up_loss)
up_G_gradients =up_G_optim.get_unscaled_gradients( tape.gradient(
scaled_gen_up_loss, up_G.trainable_variables
))
up_D_gradients = up_D_optim.get_unscaled_gradients( tape.gradient(
scaled_disc_loss, up_D.trainable_variables
))
up_G_optim.apply_gradients(
zip(up_G_gradients, up_G.trainable_variables)
)
up_D_optim.apply_gradients(
zip(up_D_gradients, up_D.trainable_variables)
)
return (
real_human,
real_anime,
fake_anime,
cycled_human,
fake_human,
cycled_anime,
same_human,
same_anime,
fake_anime_up,
real_anime_up,
gen_anime_loss,
gen_human_loss,
disc_human_loss,
disc_anime_loss,
gen_up_loss,
dis_up_loss
)
images[i] = Variable(images[i]).cuda()
text_len = Variable(text_len).cuda()
#print(text.shape, text_len)
hidden = text_encoder.init_hidden(2)
words_embs, sent_embs = text_encoder(text, text_len, hidden)
sent_embs = sent_embs.detach()
real_labels = real_labels.detach()
fake_labels = fake_labels.detach()
fake_images = g_net(noise, sent_embs)
for i in range(len(d_nets)):
d_nets[i].zero_grad()
errD = discriminator_loss(d_nets[i], images[i].detach(), fake_images[i].detach(), sent_embs, real_labels.detach(), fake_labels.detach())
errD.backward()
optimizersD[i].step()
total_error_d += errD
g_net.zero_grad()
errG_total = generator_loss(d_nets, fake_images, sent_embs, real_labels)
errG_total.backward()
optimizerG.step()
total_error_g += errG_total
print('total error: g: ', total_error_g, ' d: ', total_error_d)
save_single_image(fake_images[2], 'fake' + str(epoch) + '.png')
torch.save(g_net, 'G_NET.pth')
for i in range(len(d_nets)):
real_logits = discriminator(images, training=True)
fake_logits = discriminator(generated_images, training=True)
real_logits_rot = discriminator(images_rot,
training=True,
predict_rotation=True)
fake_logits_rot = discriminator(generated_images_rot,
training=True,
predict_rotation=True)
if second_unpaired is True:
generated_images_2 = generator(noise_2, training=True)
fake_logits_2 = discriminator(generated_images_2,
training=True)
disc_loss_2 = discriminator_loss(real_logits, fake_logits_2,
rotation_n,
real_logits_rot) # [] CHECK
gen_loss = generator_loss(fake_logits, rotation_n, fake_logits_rot)
disc_loss = discriminator_loss(real_logits, fake_logits,
rotation_n, real_logits_rot)
gradients_of_generator = gen_tape.gradient(gen_loss,
generator.variables)
gradients_of_discriminator = disc_tape.gradient(
disc_loss, discriminator.variables)
# gradients_of_discriminator_rot = disc_rot_tape.gradient(disc_loss_rot, discriminator.variables)
generator_optimizer.apply_gradients(
zip(gradients_of_generator, generator.variables))
discriminator_optimizer.apply_gradients(
