def update_netd(self, x_train):
"""更新判别器网络,
Update D network: Ladv = |f(real) - f(fake)|_2
args:
----
x_train: 训练的样本,torch.Tensor,[batch,channel,h,w],
return:
----
latent_i:从生成网络编码器中出来的向量,torch.Tensor,
latent_o:从生成网络重构编码器中出来的向量,torch.Tensor,
err_d_real: Ladv,
err_d_fake: Ladv.
"""
current_batchsize = x_train.shape[0]
label = torch.empty(size = (current_batchsize, ), dtype = torch.float32, device = self.device)
self.netd.zero_grad()
#train with real
label.data.resize_(current_batchsize).fill_(self.real_label) #用1填满便签张量
_, feat_real = self.netd(x_train)
#train with fake
label.data.resize_(current_batchsize).fill_(self.fake_label) #用0填满便签张量
self.fake, latent_i, latent_o = self.netg(x_train)
_, feat_fake = self.netd(self.fake.detach())
err_d = l2_loss(feat_real, feat_fake)
err_d_real = err_d
err_d_fake = err_d
err_d.backward()
self.optimizer_d.step()
return latent_i, latent_o, err_d_real, err_d_fake
def update_core(self):
#convert incoming array into variables with either cpu/gpu compatibility
data = Variable(self.converter(self.get_iterator('main').next(), self.device))
n, c, h, w = data.shape
# Get the ground truth and the sequential inout that is to be fed to the
# network
seq, gt = split_axis(data, [c-3], 1)
# get rid of memory
del data
output = None
total_loss_dis_adv = 0
total_loss_gen_adv = 0
for i in range(1, 5):
# Downscaling of ground truth images for loss calculations
if i != 4:
downscaled_gt = resize_images(gt, (int(h / 2 ** (4 - i)),
int(w / 2 ** (4 - i))))
downscaled_seq = resize_images(seq, (int(h / 2 ** (4 - i)),
int(w / 2 ** (4 - i))))
else:
downscaled_gt = gt
downscaled_seq = seq
output = self.GenNetwork.singleforward(i, downscaled_seq,
output)
dis_output_fake = self.DisNetwork.singleforward(i,output)
dis_outplut_real = self.DisNetwork.singleforward(i, downscaled_gt)
loss_dis = (loss_target1(dis_outplut_real) + loss_target0(dis_output_fake)) / 2
loss_gen = loss_target1(dis_output_fake)
total_loss_dis_adv += loss_dis
total_loss_gen_adv += loss_gen
loss_l2 = l2_loss(output, gt)
loss_gdl = gradient_loss(output, gt)
composite_gen_loss = self.LAM_LP*loss_l2 + self.LAM_GDL*loss_gdl + self.LAM_ADV*total_loss_gen_adv
report({'L2Loss':loss_l2},self.GenNetwork)
report({'GDL':loss_gdl},self.GenNetwork)
report({'AdvLoss':total_loss_gen_adv},self.GenNetwork)
report({'DisLoss':total_loss_dis_adv},self.DisNetwork)
report({'CompositeGenLoss':composite_gen_loss},self.GenNetwork)
# TODO: Come up with a more elegant way
self.DisNetwork.cleargrads()
self.GenNetwork.cleargrads()
composite_gen_loss.backward()
self._optimizers["GeneratorNetwork"].update()
self.DisNetwork.cleargrads()
self.GenNetwork.cleargrads()
total_loss_dis_adv.backward()
self._optimizers["DiscriminatorNetwork"].update()
def train_step(inputs):
with tf.GradientTape() as gen_tape, tf.GradientTape() as dis_tape:
outputs = model(inputs)
generation_B = outputs[0]
generation_B2 = outputs[1]
generation_C = outputs[2]
generation_A = outputs[3]
cycle_A = outputs[4]
generation_A_identity = outputs[5]
generation_alpha_back = outputs[6]
discrimination_B_fake = outputs[7]
discrimination_B_real = outputs[8]
discrimination_alpha_fake = outputs[9]
discrimination_A_dot_fake = outputs[10]
discrimination_A_dot_real = outputs[11]
sr = outputs[12]
sf = outputs[13]
w1 = outputs[14]
w2 = outputs[15]
w3 = outputs[16]
w4 = outputs[17]
interpolate_identity = outputs[18]
interpolate_B = outputs[19]
interpolate_alpha = outputs[20]
# cycle loss.
cycle_loss = l1_loss(y=inputs[0], y_hat=cycle_A)
# identity loss.
identity_loss = l1_loss(y=inputs[0], y_hat=generation_A_identity)
# backward loss.
backward_loss = l1_loss(y=inputs[0], y_hat=generation_alpha_back)
# mode-seeking loss.
mode_seeking_loss = tf.divide(l1_loss(y=inputs[0], y_hat=inputs[1]),
l1_loss(y=generation_B, y_hat=generation_B2))
# triangle loss.
triangle_loss = l1_loss(y=inputs[0], y_hat=generation_A)
# generator loss.
generator_loss_A2B = l2_loss(y=tf.ones_like(discrimination_B_fake), y_hat=discrimination_B_fake)
# two-step generator loss.
two_step_generator_loss_A = l2_loss(y=tf.ones_like(discrimination_A_dot_fake), y_hat=discrimination_A_dot_fake)
# discriminator loss.
discriminator_loss_B_real = l2_loss(y=tf.ones_like(discrimination_B_real), y_hat=discrimination_B_real)
discriminator_loss_B_fake = l2_loss(y=tf.zeros_like(discrimination_B_fake), y_hat=discrimination_B_fake)
discriminator_loss_B = (discriminator_loss_B_real + discriminator_loss_B_fake) / 2
discriminator_loss_alpha = l2_loss(y=tf.zeros_like(discrimination_alpha_fake), y_hat=discrimination_alpha_fake)
# conditional adversarial loss.
discriminator_loss_cond_sr = l2_loss(y=tf.ones_like(sr), y_hat=sr)
discriminator_loss_cond_sf = l2_loss(y=tf.zeros_like(sf), y_hat=sf)
discriminator_loss_cond_w1 = l2_loss(y=tf.zeros_like(w1), y_hat=w1)
discriminator_loss_cond_w2 = l2_loss(y=tf.zeros_like(w2), y_hat=w2)
discriminator_loss_cond_w3 = l2_loss(y=tf.zeros_like(w3), y_hat=w3)
discriminator_loss_cond_w4 = l2_loss(y=tf.zeros_like(w4), y_hat=w4)
discriminator_loss_cond = discriminator_loss_cond_sr + discriminator_loss_cond_sf + discriminator_loss_cond_w1 \
+ discriminator_loss_cond_w2 + discriminator_loss_cond_w3 + discriminator_loss_cond_w4
generator_loss_cond_sf = l2_loss(y=tf.ones_like(sf), y_hat=sf)
# interpolation loss.
discriminator_loss_interp_AB = l2_loss(y=tf.zeros_like(interpolate_identity),
y_hat=interpolate_identity) if rnd == 0 else l2_loss(
y=tf.zeros_like(interpolate_B), y_hat=interpolate_B)
discriminator_loss_interp_alpha = l2_loss(
y=tf.ones_like(interpolate_alpha) * tf.cast(tf.reshape(inputs[7], [-1, 1, 1, 1]), tf.float32),
y_hat=interpolate_alpha) if rnd == 0 else l2_loss(
y=tf.ones_like(interpolate_alpha) * tf.cast(tf.reshape(1. - inputs[7], [-1, 1, 1, 1]), tf.float32),
y_hat=interpolate_alpha)
discriminator_loss_interp = discriminator_loss_interp_AB + discriminator_loss_interp_alpha
generator_loss_interp_alpha = l2_loss(y=tf.zeros_like(interpolate_alpha), y_hat=interpolate_alpha)
# merge the losses.
generator_loss = generator_loss_A2B + hp.lambda_backward * backward_loss + mode_seeking_loss + hp.lambda_cycle \
* cycle_loss + hp.lambda_identity * identity_loss + hp.lambda_triangle * triangle_loss + \
hp.lambda_conditional * generator_loss_cond_sf + hp.lambda_interp * generator_loss_interp_alpha
discriminator_loss = discriminator_loss_B + hp.lambda_interp * discriminator_loss_interp + \
hp.lambda_conditional * discriminator_loss_cond
# Optimizers
generator_vars = model.generator.trainable_variables
discriminator_vars = model.discriminator.trainable_variables + model.adversarial.trainable_variables + \
model.interpolate.trainable_variables + model.matching.trainable_variables
grad_gen = gen_tape.gradient(generator_loss, sources=generator_vars)
grad_dis = dis_tape.gradient(discriminator_loss, sources=discriminator_vars)
generator_optimizer.apply_gradients(zip(grad_gen, generator_vars))
discriminator_optimizer.apply_gradients(zip(grad_dis, discriminator_vars))
# update summaries.
gen_loss_summary.update_state(generator_loss)
dis_loss_summary.update_state(discriminator_loss)
cycle_loss_summary.update_state(cycle_loss)
identity_loss_summary.update_state(identity_loss)
triangle_loss_summary.update_state(triangle_loss)
backward_loss_summary.update_state(backward_loss)
interpolation_loss_summary.update_state(generator_loss_interp_alpha)
conditional_loss_summary.update_state(generator_loss_cond_sf)
generator_loss_A2B_summary.update_state(generator_loss_A2B)
mode_seeking_loss_summary.update_state(mode_seeking_loss)
discriminator_B_loss_summary.update_state(discriminator_loss_B)
discriminator_loss_cond_summary.update_state(discriminator_loss_cond)
discriminator_loss_interp_summary.update_state(discriminator_loss_interp)
# tipImg = paf_maps[k].to('cpu').detach().numpy().copy()
# print(tipImg.shape, tipImg.min(), tipImg.max())
# tipImg = np.resize(tipImg, (135, 240, 1))
# tipImg = (tipImg - tipImg.min())/(tipImg.max()-tipImg.min())*255
# tipImg = np.array(tipImg, dtype='uint8')
# cv2.imwrite("images/rose_slow_ann/%d.jpg"%(k), tipImg)
# tipImg = keypoint_maps[18].to('cpu').detach().numpy().copy()
# print(tipImg.shape, tipImg.min(), tipImg.max())
# tipImg = np.resize(tipImg, (135, 240, 1))
# tipImg = (tipImg - tipImg.min())/(tipImg.max()-tipImg.min())*255
# tipImg = np.array(tipImg, dtype='uint8')
# cv2.imwrite("images/rose_slow_ann/%d.jpg"%(50), tipImg)
for loss_idx in range(len(total_losses) // 2):
losses.append(
l2_loss(stages_output[loss_idx * 2][0][0], paf_maps[0],
images.shape[0]))
losses.append(
l2_loss(stages_output[loss_idx * 2][0][1], paf_maps[1],
images.shape[0]))
losses.append(
l2_loss(stages_output[loss_idx * 2 + 1][0],
keypoint_maps[0], images.shape[0]))
total_losses[loss_idx *
2] += losses[-3].item() / batches_per_iter
total_losses[loss_idx *
2] += losses[-2].item() / batches_per_iter
total_losses[loss_idx * 2 +
1] += losses[-1].item() / batches_per_iter
loss = losses[0]
for loss_idx in range(1, len(losses)):