def _optimize_generator(self,
fake_images,
prev_image,
condition,
objects,
aux_reg,
mask,
mu,
logvar,
seg_reg=0,
seg_fake=None,
seg_gt=None):
self.generator.zero_grad()
d_fake, aux_fake, _ = self.discriminator(fake_images, condition,
prev_image)
g_loss = self._generator_masked_loss(d_fake, aux_fake, objects,
aux_reg, mu, logvar, mask,
seg_reg, seg_fake, seg_gt)
g_loss.backward(retain_graph=True)
gen_grad_norm = _recurrent_gan.get_grad_norm(
self.generator.parameters())
self.generator_optimizer.step()
g_loss_scalar = g_loss.item()
gen_grad_norm_scalar = gen_grad_norm.item()
del g_loss
del gen_grad_norm
gc.collect()
return g_loss_scalar, gen_grad_norm_scalar
def _optimize_rnn(self):
torch.nn.utils.clip_grad_norm_(self.rnn.parameters(), self.cfg.grad_clip)
rnn_grad_norm = _recurrent_gan.get_grad_norm(self.rnn.parameters())
self.rnn_optimizer.step()
self.rnn.zero_grad()
gru_grad_norm = None
torch.nn.utils.clip_grad_norm_(self.sentence_encoder.parameters(), self.cfg.grad_clip)
gru_grad_norm = _recurrent_gan.get_grad_norm(self.sentence_encoder.parameters())
self.sentence_encoder_optimizer.step()
self.sentence_encoder.zero_grad()
ce_grad_norm = _recurrent_gan.get_grad_norm(self.condition_encoder.parameters())
ie_grad_norm = _recurrent_gan.get_grad_norm(self.image_encoder.parameters())
self.feature_encoders_optimizer.step()
self.condition_encoder.zero_grad()
self.image_encoder.zero_grad()
return rnn_grad_norm, gru_grad_norm, ce_grad_norm, ie_grad_norm
def _optimize_discriminator(self,
real_images,
fake_images,
prev_image,
condition,
mask,
objects,
gp_reg=0,
aux_reg=0):
"""Discriminator is updated every step independent of batch_size
RNN and the generator
"""
wrong_images = torch.cat((real_images[1:], real_images[0:1]), dim=0)
wrong_prev = torch.cat((prev_image[1:], prev_image[0:1]), dim=0)
self.discriminator.zero_grad()
real_images.requires_grad_()
d_real, aux_real, _ = self.discriminator(real_images, condition,
prev_image)
d_fake, aux_fake, _ = self.discriminator(fake_images, condition,
prev_image)
d_wrong, _, _ = self.discriminator(wrong_images, condition, wrong_prev)
d_loss, aux_loss = self._discriminator_masked_loss(
d_real, d_fake, d_wrong, aux_real, aux_fake, objects, aux_reg,
mask)
d_loss.backward(retain_graph=True)
if gp_reg:
reg = gp_reg * self._masked_gradient_penalty(
d_real, real_images, mask)
reg.backward(retain_graph=True)
grad_norm = _recurrent_gan.get_grad_norm(
self.discriminator.parameters())
self.discriminator_optimizer.step()
d_loss_scalar = d_loss.item()
d_real_np = d_real.cpu().data.numpy()
d_fake_np = d_fake.cpu().data.numpy()
aux_loss_scalar = aux_loss.item() if isinstance(
aux_loss, torch.Tensor) else aux_loss
grad_norm_scalar = grad_norm.item()
del d_loss
del d_real
del d_fake
del aux_loss
del grad_norm
gc.collect()
return d_loss_scalar, d_real_np, d_fake_np, aux_loss_scalar, grad_norm_scalar